Merge pull request #17 from fakerybakery/main

Make ACE-Step `pip`-installable package
2025-05-07 23:34:34 +08:00
parent 9bf891fb28 df2b9d5f77
commit be5c795257
44 changed files with 3078 additions and 1405 deletions
@@ -176,6 +176,7 @@ cython_debug/
 !requirements.txt
 *.log
 *.flac
+*.wav
 minio_config.yaml
 .history/*
 __pycache__/*
@@ -200,3 +201,4 @@ app_demo.py
 ui/components_demo.py
 data_sampler_demo.py
 pipeline_ace_step_demo.py
+*.wav
@@ -34,7 +34,7 @@ WORKDIR /app
 # Clone the repository
 RUN git clone https://github.com/ace-step/ACE-Step.git .

-# Install specific PyTorch version compatible with CUDA 12.2
+# Install specific PyTorch version compatible with CUDA 12.6
 RUN pip3 install --no-cache-dir --upgrade pip && \
    pip3 install --no-cache-dir hf_transfer peft && \
    pip3 install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu126
@@ -55,4 +55,4 @@ HEALTHCHECK --interval=60s --timeout=10s --start-period=5s --retries=5 \
  CMD curl -f http://localhost:7865/ || exit 1

 # Command to run the application with GPU support
-CMD ["python3", "app.py", "--server_name", "0.0.0.0", "--bf16", "true"]
+CMD ["python3", "acestep/gui.py", "--server_name", "0.0.0.0", "--bf16", "true"]
@@ -1,22 +1,21 @@
-<h1 align="center">ACE-Step</h1>
-<h1 align="center">A Step Towards Music Generation Foundation Model</h1>
-<p align="center">
-    <a href="https://ace-step.github.io/">Project</a> |
-    <a href="https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B">Hugging Face</a> |
-    <a href="https://modelscope.cn/models/ACE-Step/ACE-Step-v1-3.5B">ModelScope</a> |
-    <a href="https://huggingface.co/spaces/ACE-Step/ACE-Step">Space Demo</a> |
-     <a href="https://discord.gg/rjAZz2xBdG">Discord</a> 
-</p>
+# ACE-Step
+
+A Step Towards Music Generation Foundation Model
+
+[Project Page](https://ace-step.github.io/) |
+[Checkpoints](https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B) |
+[Space Demo](https://huggingface.co/spaces/ACE-Step/ACE-Step) |
+[Discord](https://discord.gg/rjAZz2xBdG) 

 ---
 <p align="center">
-    <img src="./fig/orgnization_logos.png" width="100%" alt="StepFun Logo">
+    <img src="./assets/orgnization_logos.png" width="100%" alt="StepFun Logo">
 </p>

 ## Table of Contents

- [Features](#-features)
 - [Installation](#-installation)
+- [Features](#-features)
 - [Usage](#-usage)
 - [Train](#-train)

@@ -24,21 +23,43 @@

 - 🚀 2025.05.06: Open source demo code and model

-## TODOs📋
+## 📋 Roadmap
+
 - [x] Release training code 🔥
 - [x] Release LoRA training code 🔥
- [ ] Release RapMachine lora 🎤
+- [ ] Release RapMachine LoRA 🎤
 - [ ] Release ControlNet training code 🔥
- [ ] Release Singing2Accompaniment controlnet 🎮
- [ ] Release evaluation performance and technical report  📄
+- [ ] Release Singing2Accompaniment ControlNet 🎮
+- [ ] Release evaluation performance and technical report 📄
+
+## 📦 Installation
+
+### Quick Start
+
+```bash
+pip install git+https://github.com/ace-step/ACE-Step
+```
+
+On Windows, you may need to run the following before running the command above:
+
+```bash
+pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
+```
+
+Then, launch the GUI:
+
+```bash
+acestep # Launch the GUI
+```
+
+The GUI works on Windows, macOS, and Linux.

 ## 🏗️ Architecture

 <p align="center">
-    <img src="./fig/ACE-Step_framework.png" width="100%" alt="ACE-Step Framework">
+    <img src="./assets/ACE-Step_framework.png" width="100%" alt="ACE-Step Framework">
 </p>

-
 ## 📝 Abstract

 We introduce ACE-Step, a novel open-source foundation model for music generation that overcomes key limitations of existing approaches and achieves state-of-the-art performance through a holistic architectural design. Current methods face inherent trade-offs between generation speed, musical coherence, and controllability. For instance, LLM-based models (e.g., Yue, SongGen) excel at lyric alignment but suffer from slow inference and structural artifacts. Diffusion models (e.g., DiffRhythm), on the other hand, enable faster synthesis but often lack long-range structural coherence.
@@ -50,42 +71,49 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i
 ## ✨ Features

 <p align="center">
-    <img src="./fig/application_map.png" width="100%" alt="ACE-Step Framework">
+    <img src="./assets/application_map.png" width="100%" alt="ACE-Step Framework">
 </p>

 ### 🎯 Baseline Quality

 #### 🌈 Diverse Styles & Genres
+
 - 🎸 Supports all mainstream music styles with various description formats including short tags, descriptive text, or use-case scenarios
 - 🎷 Capable of generating music across different genres with appropriate instrumentation and style

 #### 🌍 Multiple Languages
+
 - 🗣️ Supports 19 languages with top 10 well-performing languages including:
  - 🇺🇸 English, 🇨🇳 Chinese, 🇷🇺 Russian, 🇪🇸 Spanish, 🇯🇵 Japanese, 🇩🇪 German, 🇫🇷 French, 🇵🇹 Portuguese, 🇮🇹 Italian, 🇰🇷 Korean
 - ⚠️ Due to data imbalance, less common languages may underperform

 #### 🎻 Instrumental Styles
+
 - 🎹 Supports various instrumental music generation across different genres and styles
 - 🎺 Capable of producing realistic instrumental tracks with appropriate timbre and expression for each instrument
 - 🎼 Can generate complex arrangements with multiple instruments while maintaining musical coherence

 #### 🎤 Vocal Techniques
+
 - 🎙️ Capable of rendering various vocal styles and techniques with good quality
 - 🗣️ Supports different vocal expressions including various singing techniques and styles

 ### 🎛️ Controllability

 #### 🔄 Variations Generation
+
 - ⚙️ Implemented using training-free, inference-time optimization techniques
 - 🌊 Flow-matching model generates initial noise, then uses trigFlow's noise formula to add additional Gaussian noise
 - 🎚️ Adjustable mixing ratio between original initial noise and new Gaussian noise to control variation degree

 #### 🎨 Repainting
+
 - 🖌️ Implemented by adding noise to the target audio input and applying mask constraints during the ODE process
 - 🔍 When input conditions change from the original generation, only specific aspects can be modified while preserving the rest
 - 🔀 Can be combined with Variations Generation techniques to create localized variations in style, lyrics, or vocals

 #### ✏️ Lyric Editing
+
 - 💡 Innovatively applies flow-edit technology to enable localized lyric modifications while preserving melody, vocals, and accompaniment
 - 🔄 Works with both generated content and uploaded audio, greatly enhancing creative possibilities
 - ℹ️ Current limitation: can only modify small segments of lyrics at once to avoid distortion, but multiple edits can be applied sequentially
@@ -93,11 +121,13 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i
 ### 🚀 Applications

 #### 🎤 Lyric2Vocal (LoRA)
+
 - 🔊 Based on a LoRA fine-tuned on pure vocal data, allowing direct generation of vocal samples from lyrics
 - 🛠️ Offers numerous practical applications such as vocal demos, guide tracks, songwriting assistance, and vocal arrangement experimentation
 - ⏱️ Provides a quick way to test how lyrics might sound when sung, helping songwriters iterate faster

 #### 📝 Text2Samples (LoRA)
+
 - 🎛️ Similar to Lyric2Vocal, but fine-tuned on pure instrumental and sample data
 - 🎵 Capable of generating conceptual music production samples from text descriptions
 - 🧰 Useful for quickly creating instrument loops, sound effects, and musical elements for production
@@ -105,16 +135,19 @@ Rather than building yet another end-to-end text-to-music pipeline, our vision i
 ### 🔮 Coming Soon

 #### 🎤 RapMachine
+
 - 🔥 Fine-tuned on pure rap data to create an AI system specialized in rap generation
 - 🏆 Expected capabilities include AI rap battles and narrative expression through rap
 - 📚 Rap has exceptional storytelling and expressive capabilities, offering extraordinary application potential

 #### 🎛️ StemGen
+
 - 🎚️ A controlnet-lora trained on multi-track data to generate individual instrument stems
 - 🎯 Takes a reference track and specified instrument (or instrument reference audio) as input
 - 🎹 Outputs an instrument stem that complements the reference track, such as creating a piano accompaniment for a flute melody or adding jazz drums to a lead guitar

 #### 🎤 Singing2Accompaniment
+
 - 🔄 The reverse process of StemGen, generating a mixed master track from a single vocal track
 - 🎵 Takes a vocal track and specified style as input to produce a complete vocal accompaniment
 - 🎸 Creates full instrumental backing that complements the input vocals, making it easy to add professional-sounding accompaniment to any vocal recording
@@ -134,89 +167,40 @@ We have evaluated ACE-Step across different hardware setups, yielding the follow

 We use RTF (Real-Time Factor) to measure the performance of ACE-Step. Higher values indicate faster generation speed. 27.27x means to generate 1 minute of music, it takes 2.2 seconds (60/27.27). The performance is measured on a single GPU with batch size 1 and 27 steps.

-
-## 💻 Installation
-
-### Prerequisites
-
-* Make sure you have Python installed. You can download it from [python.org](https://www.python.org/).
-* You will also need either `Conda` (recommended) or `venv`.
-
-### Environment Setup
-
-It is highly recommended to use a virtual environment to manage project dependencies and avoid conflicts. Choose **one** of the following methods (Conda or venv):
-
-#### Option 1: Using Conda
-
-1.  **Create the environment** named `ace_step` with Python 3.10:
-    ```bash
-    conda create -n ace_step python=3.10 -y
-    ```
-
-2.  **Activate the environment:**
-    ```bash
-    conda activate ace_step
-    ```
-
-#### Option 2: Using venv
-
-1.  **Ensure you are using the correct Python version.**
-
-2.  **Create the virtual environment** (commonly named `venv`):
-    ```bash
-    python -m venv venv 
-    ```
-
-3.  **Activate the environment:**
-    * **On Windows (cmd.exe):**
-        ```bash
-        venv\Scripts\activate.bat
-        ```
-    * **On Windows (PowerShell):**
-        ```powershell
-        .\venv\Scripts\Activate.ps1 
-        ```
-        *(If you encounter execution policy errors, you might need to run `Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope Process` first)*
-    * **On Linux / macOS (bash/zsh):**
-        ```bash
-        source venv/bin/activate
-        ```
-
-3.  **Install dependencies** from the `requirements.txt` file:
-
-    for macOS/Linux users:
-    ```bash
-    pip install -r requirements.txt
-    ```
-
-    for Windows users:
-    ```bash
-    # Install PyTorch, TorchAudio, and TorchVision for Windows
-    # replace cu126 with your CUDA version
-    # replace torchvision and torchaudio with your version
-    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
-
-    # then install other dependencies
-    pip install -r requirements.txt
-    ```
-
 ## 🚀 Usage

-![Demo Interface](fig/demo_interface.png)
+![Demo Interface](assets/demo_interface.png)

 ### 🔍 Basic Usage

 ```bash
-python app.py
+acestep
 ```

 ### ⚙️ Advanced Usage

 ```bash
-python app.py --checkpoint_path /path/to/checkpoint --port 7865 --device_id 0 --share true --bf16 true
+acestep --checkpoint_path /path/to/checkpoint --port 7865 --device_id 0 --share true --bf16 true
 ```

-If you are using MacOS, please use `--bf16 false` to avoid errors.
+If you are using macOS, please use `--bf16 false` to avoid errors.
+
+#### 🔍 API Usage
+
+```python
+from acestep.api import ACEStep
+
+model = ACEStep()
+
+audio_out = model.infer(
+    prompt="upbeat pop, catchy melody, female singer",
+    lyrics="[verse]\nSun is shining bright today\nFeeling happy, come what may",
+    audio_duration=5.0, # 5 seconds
+    infer_step=20        # Fewer steps for speed
+)
+
+audio_out.save_wav("output.wav")
+```

 #### 🛠️ Command Line Arguments

@@ -268,11 +252,11 @@ The ACE-Step interface provides several tabs for different music generation and
 - 📐 Specify left and right extension lengths
 - 🔍 Choose the source audio to extend

-## Examples
+## 📂 Examples

 The `examples/input_params` directory contains sample input parameters that can be used as references for generating music.

-## Train
+## 🔨 Train

 ### Prerequisites
 1. Prepare the environment as described in the installation section.
@@ -293,17 +277,17 @@ The `examples/input_params` directory contains sample input parameters that can
 Example dataset entry:
 ```json
 {
-  "keys": "1ce52937-cd1d-456f-967d-0f1072fcbb58",
-  "tags": ["pop", "acoustic", "ballad", "romantic", "emotional"],
-  "speaker_emb_path": "",
-  "norm_lyrics": "I love you, I love you, I love you",
-  "recaption": {
-    "simplified": "pop",
-    "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.",
-    "use_cases": "Suitable for background music in romantic films or during intimate moments.",
-    "analysis": "pop, ballad, piano, guitar, slow tempo, romantic, emotional"
-  }
+	"keys": "1ce52937-cd1d-456f-967d-0f1072fcbb58",
+	"tags": ["pop", "acoustic", "ballad", "romantic", "emotional"],
+	"speaker_emb_path": "",
+	"norm_lyrics": "I love you, I love you, I love you",
+	"recaption": {
+		"simplified": "pop",
+		"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.",
+		"use_cases": "Suitable for background music in romantic films or during intimate moments.",
+		"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):
 ```json
 {
-  "r": 16,
-  "lora_alpha": 32,
-  "target_modules": [
-    "speaker_embedder",
-    "linear_q",
-    "linear_k",
-    "linear_v",
-    "to_q",
-    "to_k",
-    "to_v",
-    "to_out.0"
-  ]
+	"r": 16,
+	"lora_alpha": 32,
+	"target_modules": [
+		"speaker_embedder",
+		"linear_q",
+		"linear_k",
+		"linear_v",
+		"to_q",
+		"to_k",
+		"to_v",
+		"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)
 - `--val_check_interval`: Validation check interval (default: None)

-## 📜 License&Disclaimer
+## 📜 License & Disclaimer

 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
@misc{gong2025acestep,
-  title={ACE-Step: A Step Towards Music Generation Foundation Model},
-  author={Junmin Gong, Wenxiao Zhao, Sen Wang, Shengyuan Xu, Jing Guo}, 
-  howpublished={\url{https://github.com/ace-step/ACE-Step}},
-  year={2025},
-  note={GitHub repository}
+	title={ACE-Step: A Step Towards Music Generation Foundation Model},
+	author={Junmin Gong, Wenxiao Zhao, Sen Wang, Shengyuan Xu, Jing Guo}, 
+	howpublished={\url{https://github.com/ace-step/ACE-Step}},
+	year={2025},
+	note={GitHub repository}
 }
 ```
@@ -0,0 +1,7 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
@@ -25,7 +25,9 @@ def project(
    v1 = torch.nn.functional.normalize(v1, dim=dims)
    v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
    v0_orthogonal = v0 - v0_parallel
-    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)
+    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(
+        device_type
+    )


 def apg_forward(
@@ -65,7 +67,11 @@ def cfg_double_condition_forward(
    guidance_scale_text,
    guidance_scale_lyric,
 ):
-    return (1 - guidance_scale_text) * uncond_output + (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output + guidance_scale_lyric * cond_output 
+    return (
+        (1 - guidance_scale_text) * uncond_output
+        + (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output
+        + guidance_scale_lyric * cond_output
+    )


 def optimized_scale(positive_flat, negative_flat):
@@ -74,7 +80,7 @@ def optimized_scale(positive_flat, negative_flat):
    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)

    # Squared norm of uncondition
-    squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
+    squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8

    # st_star = v_cond^T * v_uncond / ||v_uncond||^2
    st_star = dot_product / squared_norm
@@ -82,14 +88,23 @@ def optimized_scale(positive_flat, negative_flat):
    return st_star


-def cfg_zero_star(noise_pred_with_cond, noise_pred_uncond, guidance_scale, i, zero_steps=1, use_zero_init=True):
+def cfg_zero_star(
+    noise_pred_with_cond,
+    noise_pred_uncond,
+    guidance_scale,
+    i,
+    zero_steps=1,
+    use_zero_init=True,
+):
    bsz = noise_pred_with_cond.shape[0]
    positive_flat = noise_pred_with_cond.view(bsz, -1)
    negative_flat = noise_pred_uncond.view(bsz, -1)
    alpha = optimized_scale(positive_flat, negative_flat)
    alpha = alpha.view(bsz, 1, 1, 1)
    if (i <= zero_steps) and use_zero_init:
-        noise_pred = noise_pred_with_cond * 0.
+        noise_pred = noise_pred_with_cond * 0.0
    else:
-        noise_pred = noise_pred_uncond * alpha + guidance_scale * (noise_pred_with_cond - noise_pred_uncond * alpha)
+        noise_pred = noise_pred_uncond * alpha + guidance_scale * (
+            noise_pred_with_cond - noise_pred_uncond * alpha
+        )
    return noise_pred
@@ -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()
@@ -0,0 +1,9 @@
+from acestep.language_segmentation.LangSegment import LangSegment
+
+
+# release
+__version__ = "0.3.5"
+
+
+# develop
+__develop__ = "dev-0.0.1"
@@ -0,0 +1 @@
+# This file intentionally left blank for Python to recognize the directory as a package.
@@ -21,19 +21,21 @@ import re
 from collections import OrderedDict
 from typing import List

-DIGITS = {str(i): tran for i, tran in enumerate('零一二三四五六七八九')}
-UNITS = OrderedDict({
-    1: '十',
-    2: '百',
-    3: '千',
-    4: '万',
-    8: '亿',
-})
+DIGITS = {str(i): tran for i, tran in enumerate("零一二三四五六七八九")}
+UNITS = OrderedDict(
+    {
+        1: "十",
+        2: "百",
+        3: "千",
+        4: "万",
+        8: "亿",
+    }
+)

-COM_QUANTIFIERS = '(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)'
+COM_QUANTIFIERS = "(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)"

 # 分数表达式
-RE_FRAC = re.compile(r'(-?)(\d+)/(\d+)')
+RE_FRAC = re.compile(r"(-?)(\d+)/(\d+)")


 def replace_frac(match) -> str:
@@ -54,7 +56,7 @@ def replace_frac(match) -> str:


 # 百分数表达式
-RE_PERCENTAGE = re.compile(r'(-?)(\d+(\.\d+)?)%')
+RE_PERCENTAGE = re.compile(r"(-?)(\d+(\.\d+)?)%")


 def replace_percentage(match) -> str:
@@ -74,7 +76,7 @@ def replace_percentage(match) -> str:

 # 整数表达式
 # 带负号的整数 -10
-RE_INTEGER = re.compile(r'(-)' r'(\d+)')
+RE_INTEGER = re.compile(r"(-)" r"(\d+)")


 def replace_negative_num(match) -> str:
@@ -94,7 +96,7 @@ def replace_negative_num(match) -> str:

 # 编号-无符号整形
 # 00078
-RE_DEFAULT_NUM = re.compile(r'\d{3}\d*')
+RE_DEFAULT_NUM = re.compile(r"\d{3}\d*")


 def replace_default_num(match):
@@ -112,15 +114,11 @@ def replace_default_num(match):
 # RE_ASMD = re.compile(
 #     r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))')
 RE_ASMD = re.compile(
-    r'((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))')
+    r"((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))"
+)
+
+asmd_map = {"+": "加", "-": "减", "×": "乘", "÷": "除", "=": "等于"}

-asmd_map = {
-    '+': '加',
-    '-': '减',
-    '×': '乘',
-    '÷': '除',
-    '=': '等于'
-}

 def replace_asmd(match) -> str:
    """
@@ -134,24 +132,25 @@ def replace_asmd(match) -> str:


 # 次方专项
-RE_POWER = re.compile(r'[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+')
+RE_POWER = re.compile(r"[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+")

 power_map = {
-    '⁰': '0',
-    '¹': '1',
-    '²': '2',
-    '³': '3',
-    '⁴': '4',
-    '⁵': '5',
-    '⁶': '6',
-    '⁷': '7',
-    '⁸': '8',
-    '⁹': '9',
-    'ˣ': 'x',
-    'ʸ': 'y',
-    'ⁿ': 'n'
+    "⁰": "0",
+    "¹": "1",
+    "²": "2",
+    "³": "3",
+    "⁴": "4",
+    "⁵": "5",
+    "⁶": "6",
+    "⁷": "7",
+    "⁸": "8",
+    "⁹": "9",
+    "ˣ": "x",
+    "ʸ": "y",
+    "ⁿ": "n",
 }

+
 def replace_power(match) -> str:
    """
    Args:
@@ -168,10 +167,10 @@ def replace_power(match) -> str:

 # 数字表达式
 # 纯小数
-RE_DECIMAL_NUM = re.compile(r'(-?)((\d+)(\.\d+))' r'|(\.(\d+))')
+RE_DECIMAL_NUM = re.compile(r"(-?)((\d+)(\.\d+))" r"|(\.(\d+))")
 # 正整数 + 量词
 RE_POSITIVE_QUANTIFIERS = re.compile(r"(\d+)([多余几\+])?" + COM_QUANTIFIERS)
-RE_NUMBER = re.compile(r'(-?)((\d+)(\.\d+)?)' r'|(\.(\d+))')
+RE_NUMBER = re.compile(r"(-?)((\d+)(\.\d+)?)" r"|(\.(\d+))")


 def replace_positive_quantifier(match) -> str:
@@ -221,7 +220,9 @@ RE_RANGE = re.compile(
    [-~]                   # 匹配范围分隔符
    ((-?)((\d+)(\.\d+)?))  # 匹配范围结束的负数或正数（整数或小数）
    (?![\d\+\-\×÷=])       # 使用正向前瞻以确保数字范围之后没有其他数字和操作符
-    """, re.VERBOSE)
+    """,
+    re.VERBOSE,
+)


 def replace_range(match) -> str:
@@ -240,7 +241,9 @@ def replace_range(match) -> str:

 # ~至表达式
 RE_TO_RANGE = re.compile(
-    r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)')
+    r"((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)"
+)
+

 def replace_to_range(match) -> str:
    """
@@ -249,58 +252,61 @@ def replace_to_range(match) -> str:
    Returns:
        str
    """
-    result = match.group(0).replace('~', '至')
+    result = match.group(0).replace("~", "至")
    return result


-def _get_value(value_string: str, use_zero: bool=True) -> List[str]:
-    stripped = value_string.lstrip('0')
+def _get_value(value_string: str, use_zero: bool = True) -> List[str]:
+    stripped = value_string.lstrip("0")
    if len(stripped) == 0:
        return []
    elif len(stripped) == 1:
        if use_zero and len(stripped) < len(value_string):
-            return [DIGITS['0'], DIGITS[stripped]]
+            return [DIGITS["0"], DIGITS[stripped]]
        else:
            return [DIGITS[stripped]]
    else:
        largest_unit = next(
-            power for power in reversed(UNITS.keys()) if power < len(stripped))
+            power for power in reversed(UNITS.keys()) if power < len(stripped)
+        )
        first_part = value_string[:-largest_unit]
        second_part = value_string[-largest_unit:]
-        return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(
-            second_part)
+        return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(second_part)


 def verbalize_cardinal(value_string: str) -> str:
    if not value_string:
-        return ''
+        return ""

    # 000 -> '零' , 0 -> '零'
-    value_string = value_string.lstrip('0')
+    value_string = value_string.lstrip("0")
    if len(value_string) == 0:
-        return DIGITS['0']
+        return DIGITS["0"]

    result_symbols = _get_value(value_string)
    # verbalized number starting with '一十*' is abbreviated as `十*`
-    if len(result_symbols) >= 2 and result_symbols[0] == DIGITS[
-            '1'] and result_symbols[1] == UNITS[1]:
+    if (
+        len(result_symbols) >= 2
+        and result_symbols[0] == DIGITS["1"]
+        and result_symbols[1] == UNITS[1]
+    ):
        result_symbols = result_symbols[1:]
-    return ''.join(result_symbols)
+    return "".join(result_symbols)


 def verbalize_digit(value_string: str, alt_one=False) -> str:
    result_symbols = [DIGITS[digit] for digit in value_string]
-    result = ''.join(result_symbols)
+    result = "".join(result_symbols)
    if alt_one:
        result = result.replace("一", "幺")
    return result


 def num2str(value_string: str) -> str:
-    integer_decimal = value_string.split('.')
+    integer_decimal = value_string.split(".")
    if len(integer_decimal) == 1:
        integer = integer_decimal[0]
-        decimal = ''
+        decimal = ""
    elif len(integer_decimal) == 2:
        integer, decimal = integer_decimal
    else:
@@ -310,12 +316,12 @@ def num2str(value_string: str) -> str:

    result = verbalize_cardinal(integer)

-    decimal = decimal.rstrip('0')
+    decimal = decimal.rstrip("0")
    if decimal:
        # '.22' is verbalized as '零点二二'
        # '3.20' is verbalized as '三点二
        result = result if result else "零"
-        result += '点' + verbalize_digit(decimal)
+        result += "点" + verbalize_digit(decimal)
    return result


@@ -31,9 +31,15 @@ from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder

 def cross_norm(hidden_states, controlnet_input):
    # input N x T x c
-    mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)
-    mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True)
-    controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states
+    mean_hidden_states, std_hidden_states = hidden_states.mean(
+        dim=(1, 2), keepdim=True
+    ), hidden_states.std(dim=(1, 2), keepdim=True)
+    mean_controlnet_input, std_controlnet_input = controlnet_input.mean(
+        dim=(1, 2), keepdim=True
+    ), controlnet_input.std(dim=(1, 2), keepdim=True)
+    controlnet_input = (controlnet_input - mean_controlnet_input) * (
+        std_hidden_states / (std_controlnet_input + 1e-12)
+    ) + mean_hidden_states
    return controlnet_input


@@ -45,17 +51,27 @@ class Qwen2RotaryEmbedding(nn.Module):
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
-        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+        inv_freq = 1.0 / (
+            self.base
+            ** (
+                torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
+                / self.dim
+            )
+        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
-            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+            seq_len=max_position_embeddings,
+            device=self.inv_freq.device,
+            dtype=torch.get_default_dtype(),
        )

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
-        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+        t = torch.arange(
+            self.max_seq_len_cached, device=device, dtype=torch.int64
+        ).type_as(self.inv_freq)

        freqs = torch.outer(t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
@@ -82,8 +98,12 @@ class T2IFinalLayer(nn.Module):
    def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):
        super().__init__()
        self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True)
-        self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
+        self.linear = nn.Linear(
+            hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True
+        )
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, hidden_size) / hidden_size**0.5
+        )
        self.out_channels = out_channels
        self.patch_size = patch_size

@@ -95,14 +115,28 @@ class T2IFinalLayer(nn.Module):
        # 4 unpatchify
        new_height, new_width = 1, hidden_states.size(1)
        hidden_states = hidden_states.reshape(
-            shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)
+            shape=(
+                hidden_states.shape[0],
+                new_height,
+                new_width,
+                self.patch_size[0],
+                self.patch_size[1],
+                self.out_channels,
+            )
        ).contiguous()
        hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
        output = hidden_states.reshape(
-            shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])
+            shape=(
+                hidden_states.shape[0],
+                self.out_channels,
+                new_height * self.patch_size[0],
+                new_width * self.patch_size[1],
+            )
        ).contiguous()
        if width > new_width:
-            output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
+            output = torch.nn.functional.pad(
+                output, (0, width - new_width, 0, 0), "constant", 0
+            )
        elif width < new_width:
            output = output[:, :, :, :width]
        return output
@@ -131,9 +165,25 @@ class PatchEmbed(nn.Module):
        super().__init__()
        patch_size_h, patch_size_w = patch_size
        self.early_conv_layers = nn.Sequential(
-            nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias),
-            torch.nn.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True),
-            nn.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias)
+            nn.Conv2d(
+                in_channels,
+                in_channels * 256,
+                kernel_size=patch_size,
+                stride=patch_size,
+                padding=0,
+                bias=bias,
+            ),
+            torch.nn.GroupNorm(
+                num_groups=32, num_channels=in_channels * 256, eps=1e-6, affine=True
+            ),
+            nn.Conv2d(
+                in_channels * 256,
+                embed_dim,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias=bias,
+            ),
        )
        self.patch_size = patch_size
        self.height, self.width = height // patch_size_h, width // patch_size_w
@@ -153,7 +203,9 @@ class Transformer2DModelOutput(BaseOutput):
    proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None


-class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+class ACEStepTransformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
+):
    _supports_gradient_checkpointing = True

    @register_to_config
@@ -217,9 +269,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
        )
        self.num_layers = num_layers

-        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
-        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim)
-        self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True))
+        self.time_proj = Timesteps(
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=self.inner_dim
+        )
+        self.t_block = nn.Sequential(
+            nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True)
+        )

        # speaker
        self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim)
@@ -229,25 +287,30 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO

        # lyric
        self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size)
-        self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0)
+        self.lyric_encoder = LyricEncoder(
+            input_size=lyric_hidden_size, static_chunk_size=0
+        )
        self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)

        projector_dim = 2 * self.inner_dim

-        self.projectors = nn.ModuleList([
-            nn.Sequential(
-                nn.Linear(self.inner_dim, projector_dim),
-                nn.SiLU(),
-                nn.Linear(projector_dim, projector_dim),
-                nn.SiLU(),
-                nn.Linear(projector_dim, ssl_dim),
-            ) for ssl_dim in ssl_latent_dims
-        ])
+        self.projectors = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(self.inner_dim, projector_dim),
+                    nn.SiLU(),
+                    nn.Linear(projector_dim, projector_dim),
+                    nn.SiLU(),
+                    nn.Linear(projector_dim, ssl_dim),
+                )
+                for ssl_dim in ssl_latent_dims
+            ]
+        )

        self.ssl_latent_dims = ssl_latent_dims
        self.ssl_encoder_depths = ssl_encoder_depths

-        self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')
+        self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction="mean")
        self.ssl_names = ssl_names

        self.proj_in = PatchEmbed(
@@ -258,11 +321,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
            bias=True,
        )

-        self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels)
+        self.final_layer = T2IFinalLayer(
+            self.inner_dim, patch_size=patch_size, out_channels=out_channels
+        )
        self.gradient_checkpointing = False

    # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
-    def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+    def enable_forward_chunking(
+        self, chunk_size: Optional[int] = None, dim: int = 0
+    ) -> None:
        """
        Sets the attention processor to use [feed forward
        chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -281,7 +348,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
        # By default chunk size is 1
        chunk_size = chunk_size or 1

-        def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
+        def fn_recursive_feed_forward(
+            module: torch.nn.Module, chunk_size: int, dim: int
+        ):
            if hasattr(module, "set_chunk_feed_forward"):
                module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)

@@ -302,7 +371,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
    ):
        # N x T x D
        lyric_embs = self.lyric_embs(lyric_token_idx)
-        prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)
+        prompt_prenet_out, _mask = self.lyric_encoder(
+            lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1
+        )
        prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
        return prompt_prenet_out

@@ -331,8 +402,17 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
            lyric_mask=lyric_mask,
        )

-        encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)
-        encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
+        encoder_hidden_states = torch.cat(
+            [
+                encoder_spk_hidden_states,
+                encoder_text_hidden_states,
+                encoder_lyric_hidden_states,
+            ],
+            dim=1,
+        )
+        encoder_hidden_mask = torch.cat(
+            [speaker_mask, text_attention_mask, lyric_mask], dim=1
+        )
        return encoder_hidden_states, encoder_hidden_mask

    def decode(
@@ -344,12 +424,16 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
        timestep: Optional[torch.Tensor],
        ssl_hidden_states: Optional[List[torch.Tensor]] = None,
        output_length: int = 0,
-        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+        block_controlnet_hidden_states: Optional[
+            Union[List[torch.Tensor], torch.Tensor]
+        ] = None,
        controlnet_scale: Union[float, torch.Tensor] = 1.0,
        return_dict: bool = True,
    ):

-        embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
+        embedded_timestep = self.timestep_embedder(
+            self.time_proj(timestep).to(dtype=hidden_states.dtype)
+        )
        temb = self.t_block(embedded_timestep)

        hidden_states = self.proj_in(hidden_states)
@@ -361,8 +445,12 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO

        inner_hidden_states = []

-        rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])
-        encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])
+        rotary_freqs_cis = self.rotary_emb(
+            hidden_states, seq_len=hidden_states.shape[1]
+        )
+        encoder_rotary_freqs_cis = self.rotary_emb(
+            encoder_hidden_states, seq_len=encoder_hidden_states.shape[1]
+        )

        for index_block, block in enumerate(self.transformer_blocks):

@@ -377,7 +465,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO

                    return custom_forward

-                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                ckpt_kwargs: Dict[str, Any] = (
+                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                )
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states=hidden_states,
@@ -406,9 +496,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
                    inner_hidden_states.append(hidden_states)

        proj_losses = []
-        if len(inner_hidden_states) > 0 and ssl_hidden_states is not None and len(ssl_hidden_states) > 0:
+        if (
+            len(inner_hidden_states) > 0
+            and ssl_hidden_states is not None
+            and len(ssl_hidden_states) > 0
+        ):

-            for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names):
+            for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(
+                inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names
+            ):
                if ssl_hidden_state is None:
                    continue
                # 1. N x T x D1 -> N x D x D2
@@ -416,9 +512,20 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
                # 3. projection loss
                bs = inner_hidden_state.shape[0]
                proj_loss = 0.0
-                for i, (z, z_tilde) in enumerate(zip(ssl_hidden_state, est_ssl_hidden_state)):
+                for i, (z, z_tilde) in enumerate(
+                    zip(ssl_hidden_state, est_ssl_hidden_state)
+                ):
                    # 2. interpolate
-                    z_tilde = F.interpolate(z_tilde.unsqueeze(0).transpose(1, 2), size=len(z), mode='linear', align_corners=False).transpose(1, 2).squeeze(0)
+                    z_tilde = (
+                        F.interpolate(
+                            z_tilde.unsqueeze(0).transpose(1, 2),
+                            size=len(z),
+                            mode="linear",
+                            align_corners=False,
+                        )
+                        .transpose(1, 2)
+                        .squeeze(0)
+                    )

                    z_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)
                    z = torch.nn.functional.normalize(z, dim=-1)
@@ -445,7 +552,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
        lyric_mask: Optional[torch.LongTensor] = None,
        timestep: Optional[torch.Tensor] = None,
        ssl_hidden_states: Optional[List[torch.Tensor]] = None,
-        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+        block_controlnet_hidden_states: Optional[
+            Union[List[torch.Tensor], torch.Tensor]
+        ] = None,
        controlnet_scale: Union[float, torch.Tensor] = 1.0,
        return_dict: bool = True,
    ):
@@ -23,10 +23,18 @@ from diffusers.models.normalization import RMSNorm

 try:
    # from .dcformer import DCMHAttention
-    from .customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
+    from .customer_attention_processor import (
+        Attention,
+        CustomLiteLAProcessor2_0,
+        CustomerAttnProcessor2_0,
+    )
 except ImportError:
    # from dcformer import DCMHAttention
-    from customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
+    from customer_attention_processor import (
+        Attention,
+        CustomLiteLAProcessor2_0,
+        CustomerAttnProcessor2_0,
+    )


 logger = logging.get_logger(__name__)
@@ -55,13 +63,16 @@ def t2i_modulate(x, shift, scale):
    return x * (1 + scale) + shift


-def get_same_padding(kernel_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
+def get_same_padding(
+    kernel_size: Union[int, Tuple[int, ...]],
+) -> Union[int, Tuple[int, ...]]:
    if isinstance(kernel_size, tuple):
        return tuple([get_same_padding(ks) for ks in kernel_size])
    else:
        assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
        return kernel_size // 2

+
 class ConvLayer(nn.Module):
    def __init__(
        self,
@@ -187,6 +198,7 @@ class LinearTransformerBlock(nn.Module):
    """
    A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
    """
+
    def __init__(
        self,
        dim,
@@ -78,12 +78,16 @@ class CustomLiteLAProcessor2_0:
        input_ndim = hidden_states.ndim
        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)
        if encoder_hidden_states is not None:
            context_input_ndim = encoder_hidden_states.ndim
            if context_input_ndim == 4:
                batch_size, channel, height, width = encoder_hidden_states.shape
-                encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+                encoder_hidden_states = encoder_hidden_states.view(
+                    batch_size, channel, height * width
+                ).transpose(1, 2)

        batch_size = hidden_states.shape[0]

@@ -94,7 +98,11 @@ class CustomLiteLAProcessor2_0:
        value = attn.to_v(hidden_states)

        # `context` projections.
-        has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
+        has_encoder_hidden_state_proj = (
+            hasattr(attn, "add_q_proj")
+            and hasattr(attn, "add_k_proj")
+            and hasattr(attn, "add_v_proj")
+        )
        if encoder_hidden_states is not None and has_encoder_hidden_state_proj:
            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
@@ -114,7 +122,11 @@ class CustomLiteLAProcessor2_0:
        head_dim = inner_dim // attn.heads

        query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
-        key = key.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1).transpose(-1, -2)
+        key = (
+            key.transpose(-1, -2)
+            .reshape(batch_size, attn.heads, head_dim, -1)
+            .transpose(-1, -2)
+        )
        value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)

        # RoPE需要 [B, H, S, D] 输入
@@ -140,17 +152,33 @@ class CustomLiteLAProcessor2_0:

        if attention_mask is not None:
            # attention_mask: [B, S] -> [B, 1, S, 1]
-            attention_mask = attention_mask[:, None, :, None].to(key.dtype)  # [B, 1, S, 1]
-            query = query * attention_mask.permute(0, 1, 3, 2)  # [B, H, S, D] * [B, 1, S, 1]
+            attention_mask = attention_mask[:, None, :, None].to(
+                key.dtype
+            )  # [B, 1, S, 1]
+            query = query * attention_mask.permute(
+                0, 1, 3, 2
+            )  # [B, H, S, D] * [B, 1, S, 1]
            if not attn.is_cross_attention:
-                key = key * attention_mask  # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
-                value = value * attention_mask.permute(0, 1, 3, 2)  # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
+                key = (
+                    key * attention_mask
+                )  # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
+                value = value * attention_mask.permute(
+                    0, 1, 3, 2
+                )  # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度

-        if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
-            encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(key.dtype)  # [B, 1, S_enc, 1]
+        if (
+            attn.is_cross_attention
+            and encoder_attention_mask is not None
+            and has_encoder_hidden_state_proj
+        ):
+            encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(
+                key.dtype
+            )  # [B, 1, S_enc, 1]
            # 此时 key: [B, h, S_enc, D], value: [B, h, D, S_enc]
            key = key * encoder_attention_mask  # [B, h, S_enc, D] * [B, 1, S_enc, 1]
-            value = value * encoder_attention_mask.permute(0, 1, 3, 2)  # [B, h, D, S_enc] * [B, 1, 1, S_enc]
+            value = value * encoder_attention_mask.permute(
+                0, 1, 3, 2
+            )  # [B, h, D, S_enc] * [B, 1, 1, S_enc]

        query = self.kernel_func(query)
        key = self.kernel_func(key)
@@ -168,16 +196,22 @@ class CustomLiteLAProcessor2_0:

        hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)

-        hidden_states = hidden_states.view(batch_size, attn.heads * head_dim, -1).permute(0, 2, 1)
+        hidden_states = hidden_states.view(
+            batch_size, attn.heads * head_dim, -1
+        ).permute(0, 2, 1)

        hidden_states = hidden_states.to(dtype)
        if encoder_hidden_states is not None:
            encoder_hidden_states = encoder_hidden_states.to(dtype)

        # Split the attention outputs.
-        if encoder_hidden_states is not None and not attn.is_cross_attention and has_encoder_hidden_state_proj:
+        if (
+            encoder_hidden_states is not None
+            and not attn.is_cross_attention
+            and has_encoder_hidden_state_proj
+        ):
            hidden_states, encoder_hidden_states = (
-                hidden_states[:, : hidden_states_len],
+                hidden_states[:, :hidden_states_len],
                hidden_states[:, hidden_states_len:],
            )

@@ -185,13 +219,22 @@ class CustomLiteLAProcessor2_0:
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)
-        if encoder_hidden_states is not None and not attn.context_pre_only and not attn.is_cross_attention and hasattr(attn, "to_add_out"):
+        if (
+            encoder_hidden_states is not None
+            and not attn.context_pre_only
+            and not attn.is_cross_attention
+            and hasattr(attn, "to_add_out")
+        ):
            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)

        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )
        if encoder_hidden_states is not None and context_input_ndim == 4:
-            encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+            encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )

        if torch.get_autocast_gpu_dtype() == torch.float16:
            hidden_states = hidden_states.clip(-65504, 65504)
@@ -208,7 +251,9 @@ class CustomerAttnProcessor2_0:

    def __init__(self):
        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+            raise ImportError(
+                "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )

    def apply_rotary_emb(
        self,
@@ -258,23 +303,35 @@ class CustomerAttnProcessor2_0:

        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+            hidden_states = hidden_states.view(
+                batch_size, channel, height * width
+            ).transpose(1, 2)

        batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
        )

-        has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
+        has_encoder_hidden_state_proj = (
+            hasattr(attn, "add_q_proj")
+            and hasattr(attn, "add_k_proj")
+            and hasattr(attn, "add_v_proj")
+        )

        if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(
+                1, 2
+            )

        query = attn.to_q(hidden_states)

        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        elif attn.norm_cross:
-            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+            encoder_hidden_states = attn.norm_encoder_hidden_states(
+                encoder_hidden_states
+            )

        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)
@@ -300,19 +357,33 @@ class CustomerAttnProcessor2_0:
            elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
                key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)

-        if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
+        if (
+            attn.is_cross_attention
+            and encoder_attention_mask is not None
+            and has_encoder_hidden_state_proj
+        ):
            # attention_mask: N x S1
            # encoder_attention_mask: N x S2
            # cross attention 整合attention_mask和encoder_attention_mask
-            combined_mask = attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
+            combined_mask = (
+                attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
+            )
            attention_mask = torch.where(combined_mask == 1, 0.0, -torch.inf)
-            attention_mask = attention_mask[:, None, :, :].expand(-1, attn.heads, -1, -1).to(query.dtype)
+            attention_mask = (
+                attention_mask[:, None, :, :]
+                .expand(-1, attn.heads, -1, -1)
+                .to(query.dtype)
+            )

        elif not attn.is_cross_attention and attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
            # scaled_dot_product_attention expects attention_mask shape to be
            # (batch, heads, source_length, target_length)
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )

        # the output of sdp = (batch, num_heads, seq_len, head_dim)
        # TODO: add support for attn.scale when we move to Torch 2.1
@@ -320,7 +391,9 @@ class CustomerAttnProcessor2_0:
            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
        )

-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            batch_size, -1, attn.heads * head_dim
+        )
        hidden_states = hidden_states.to(query.dtype)

        # linear proj
@@ -329,7 +402,9 @@ class CustomerAttnProcessor2_0:
        hidden_states = attn.to_out[1](hidden_states)

        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+            hidden_states = hidden_states.transpose(-1, -2).reshape(
+                batch_size, channel, height, width
+            )

        if attn.residual_connection:
            hidden_states = hidden_states + residual
@@ -3,16 +3,19 @@ import math
 import torch
 from torch import nn

+
 class ConvolutionModule(nn.Module):
    """ConvolutionModule in Conformer model."""

-    def __init__(self,
-                 channels: int,
-                 kernel_size: int = 15,
-                 activation: nn.Module = nn.ReLU(),
-                 norm: str = "batch_norm",
-                 causal: bool = False,
-                 bias: bool = True):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 15,
+        activation: nn.Module = nn.ReLU(),
+        norm: str = "batch_norm",
+        causal: bool = False,
+        bias: bool = True,
+    ):
        """Construct an ConvolutionModule object.
        Args:
            channels (int): The number of channels of conv layers.
@@ -51,7 +54,7 @@ class ConvolutionModule(nn.Module):
            bias=bias,
        )

-        assert norm in ['batch_norm', 'layer_norm']
+        assert norm in ["batch_norm", "layer_norm"]
        if norm == "batch_norm":
            self.use_layer_norm = False
            self.norm = nn.BatchNorm1d(channels)
@@ -95,13 +98,13 @@ class ConvolutionModule(nn.Module):

        if self.lorder > 0:
            if cache.size(2) == 0:  # cache_t == 0
-                x = nn.functional.pad(x, (self.lorder, 0), 'constant', 0.0)
+                x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0)
            else:
                assert cache.size(0) == x.size(0)  # equal batch
                assert cache.size(1) == x.size(1)  # equal channel
                x = torch.cat((cache, x), dim=2)
-            assert (x.size(2) > self.lorder)
-            new_cache = x[:, :, -self.lorder:]
+            assert x.size(2) > self.lorder
+            new_cache = x[:, :, -self.lorder :]
        else:
            # It's better we just return None if no cache is required,
            # 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

+
 class PositionwiseFeedForward(torch.nn.Module):
    """Positionwise feed forward layer.

@@ -140,11 +144,11 @@ class PositionwiseFeedForward(torch.nn.Module):
    """

    def __init__(
-            self,
-            idim: int,
-            hidden_units: int,
-            dropout_rate: float,
-            activation: torch.nn.Module = torch.nn.ReLU(),
+        self,
+        idim: int,
+        hidden_units: int,
+        dropout_rate: float,
+        activation: torch.nn.Module = torch.nn.ReLU(),
    ):
        """Construct a PositionwiseFeedForward object."""
        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))))

+
 class Swish(torch.nn.Module):
    """Construct an Swish object."""

@@ -170,6 +175,7 @@ class Swish(torch.nn.Module):
        """Return Swish activation function."""
        return x * torch.sigmoid(x)

+
 class MultiHeadedAttention(nn.Module):
    """Multi-Head Attention layer.

@@ -180,11 +186,9 @@ class MultiHeadedAttention(nn.Module):

    """

-    def __init__(self,
-                 n_head: int,
-                 n_feat: int,
-                 dropout_rate: float,
-                 key_bias: bool = True):
+    def __init__(
+        self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
+    ):
        """Construct an MultiHeadedAttention object."""
        super().__init__()
        assert n_feat % n_head == 0
@@ -229,7 +233,7 @@ class MultiHeadedAttention(nn.Module):
        self,
        value: torch.Tensor,
        scores: torch.Tensor,
-        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
    ) -> torch.Tensor:
        """Compute attention context vector.

@@ -251,19 +255,20 @@ class MultiHeadedAttention(nn.Module):
        if mask.size(2) > 0:  # time2 > 0
            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
            # For last chunk, time2 might be larger than scores.size(-1)
-            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
-            scores = scores.masked_fill(mask, -float('inf'))
+            mask = mask[:, :, :, : scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float("inf"))
            attn = torch.softmax(scores, dim=-1).masked_fill(
-                mask, 0.0)  # (batch, head, time1, time2)
+                mask, 0.0
+            )  # (batch, head, time1, time2)

        else:
            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)

        p_attn = self.dropout(attn)
        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
-        x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
-                                                 self.h * self.d_k)
-             )  # (batch, time1, d_model)
+        x = (
+            x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)
+        )  # (batch, time1, d_model)

        return self.linear_out(x)  # (batch, time1, d_model)

@@ -274,7 +279,7 @@ class MultiHeadedAttention(nn.Module):
        value: torch.Tensor,
        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
        pos_emb: torch.Tensor = torch.empty(0),
-        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute scaled dot product attention.

@@ -308,9 +313,7 @@ class MultiHeadedAttention(nn.Module):
        """
        q, k, v = self.forward_qkv(query, key, value)
        if cache.size(0) > 0:
-            key_cache, value_cache = torch.split(cache,
-                                                 cache.size(-1) // 2,
-                                                 dim=-1)
+            key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
            k = torch.cat([key_cache, k], dim=2)
            v = torch.cat([value_cache, v], dim=2)
        new_cache = torch.cat((k, v), dim=-1)
@@ -328,11 +331,9 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
        dropout_rate (float): Dropout rate.
    """

-    def __init__(self,
-                 n_head: int,
-                 n_feat: int,
-                 dropout_rate: float,
-                 key_bias: bool = True):
+    def __init__(
+        self, n_head: int, n_feat: int, dropout_rate: float, key_bias: bool = True
+    ):
        """Construct an RelPositionMultiHeadedAttention object."""
        super().__init__(n_head, n_feat, dropout_rate, key_bias)
        # linear transformation for positional encoding
@@ -355,14 +356,12 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
            torch.Tensor: Output tensor.

        """
-        zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1),
-                               device=x.device,
-                               dtype=x.dtype)
+        zero_pad = torch.zeros(
+            (x.size()[0], x.size()[1], x.size()[2], 1), device=x.device, dtype=x.dtype
+        )
        x_padded = torch.cat([zero_pad, x], dim=-1)

-        x_padded = x_padded.view(x.size()[0],
-                                 x.size()[1],
-                                 x.size(3) + 1, x.size(2))
+        x_padded = x_padded.view(x.size()[0], x.size()[1], x.size(3) + 1, x.size(2))
        x = x_padded[:, :, 1:].view_as(x)[
            :, :, :, : x.size(-1) // 2 + 1
        ]  # only keep the positions from 0 to time2
@@ -375,7 +374,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
        value: torch.Tensor,
        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
        pos_emb: torch.Tensor = torch.empty(0),
-        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+        cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
        Args:
@@ -399,9 +398,7 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
        q = q.transpose(1, 2)  # (batch, time1, head, d_k)

        if cache.size(0) > 0:
-            key_cache, value_cache = torch.split(cache,
-                                                 cache.size(-1) // 2,
-                                                 dim=-1)
+            key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1)
            k = torch.cat([key_cache, k], dim=2)
            v = torch.cat([value_cache, v], dim=2)
        # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
@@ -431,15 +428,15 @@ class RelPositionMultiHeadedAttention(MultiHeadedAttention):
            matrix_bd = self.rel_shift(matrix_bd)

        scores = (matrix_ac + matrix_bd) / math.sqrt(
-            self.d_k)  # (batch, head, time1, time2)
+            self.d_k
+        )  # (batch, head, time1, time2)

        return self.forward_attention(v, scores, mask), new_cache


-
 def subsequent_mask(
-        size: int,
-        device: torch.device = torch.device("cpu"),
+    size: int,
+    device: torch.device = torch.device("cpu"),
 ) -> torch.Tensor:
    """Create mask for subsequent steps (size, size).

@@ -474,11 +471,11 @@ def subsequent_mask(


 def subsequent_chunk_mask(
-        size: int,
-        chunk_size: int,
-        num_left_chunks: int = -1,
-        device: torch.device = torch.device("cpu"),
-    ) -> torch.Tensor:
+    size: int,
+    chunk_size: int,
+    num_left_chunks: int = -1,
+    device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
    """Create mask for subsequent steps (size, size) with chunk size,
       this is for streaming encoder

@@ -510,15 +507,18 @@ def subsequent_chunk_mask(
        ret[i, start:ending] = True
    return ret

-def add_optional_chunk_mask(xs: torch.Tensor,
-                            masks: torch.Tensor,
-                            use_dynamic_chunk: bool,
-                            use_dynamic_left_chunk: bool,
-                            decoding_chunk_size: int,
-                            static_chunk_size: int,
-                            num_decoding_left_chunks: int,
-                            enable_full_context: bool = True):
-    """ Apply optional mask for encoder.
+
+def add_optional_chunk_mask(
+    xs: torch.Tensor,
+    masks: torch.Tensor,
+    use_dynamic_chunk: bool,
+    use_dynamic_left_chunk: bool,
+    decoding_chunk_size: int,
+    static_chunk_size: int,
+    num_decoding_left_chunks: int,
+    enable_full_context: bool = True,
+):
+    """Apply optional mask for encoder.

    Args:
        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).
            # Since we use 4 times subsampling and allow up to 1s(100 frames)
            # 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
            if chunk_size > max_len // 2 and enable_full_context:
                chunk_size = max_len
@@ -565,18 +565,17 @@ def add_optional_chunk_mask(xs: torch.Tensor,
                chunk_size = chunk_size % 25 + 1
                if use_dynamic_left_chunk:
                    max_left_chunks = (max_len - 1) // chunk_size
-                    num_left_chunks = torch.randint(0, max_left_chunks,
-                                                    (1, )).item()
-        chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size,
-                                            num_left_chunks,
-                                            xs.device)  # (L, L)
+                    num_left_chunks = torch.randint(0, max_left_chunks, (1,)).item()
+        chunk_masks = subsequent_chunk_mask(
+            xs.size(1), chunk_size, num_left_chunks, xs.device
+        )  # (L, L)
        chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
        chunk_masks = masks & chunk_masks  # (B, L, L)
    elif static_chunk_size > 0:
        num_left_chunks = num_decoding_left_chunks
-        chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size,
-                                            num_left_chunks,
-                                            xs.device)  # (L, L)
+        chunk_masks = subsequent_chunk_mask(
+            xs.size(1), static_chunk_size, num_left_chunks, xs.device
+        )  # (L, L)
        chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
        chunk_masks = masks & chunk_masks  # (B, L, L)
    else:
@@ -630,7 +629,8 @@ class ConformerEncoderLayer(nn.Module):
        if self.conv_module is not None:
            self.norm_conv = nn.LayerNorm(size, eps=1e-5)  # for the CNN module
            self.norm_final = nn.LayerNorm(
-                size, eps=1e-5)  # for the final output of the block
+                size, eps=1e-5
+            )  # for the final output of the block
        self.dropout = nn.Dropout(dropout_rate)
        self.size = size
        self.normalize_before = normalize_before
@@ -671,8 +671,7 @@ class ConformerEncoderLayer(nn.Module):
            residual = x
            if self.normalize_before:
                x = self.norm_ff_macaron(x)
-            x = residual + self.ff_scale * self.dropout(
-                self.feed_forward_macaron(x))
+            x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x))
            if not self.normalize_before:
                x = self.norm_ff_macaron(x)

@@ -680,8 +679,7 @@ class ConformerEncoderLayer(nn.Module):
        residual = x
        if self.normalize_before:
            x = self.norm_mha(x)
-        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
-                                              att_cache)
+        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache)
        x = residual + self.dropout(x_att)
        if not self.normalize_before:
            x = self.norm_mha(x)
@@ -714,7 +712,6 @@ class ConformerEncoderLayer(nn.Module):
        return x, mask, new_att_cache, new_cnn_cache


-
 class EspnetRelPositionalEncoding(torch.nn.Module):
    """Relative positional encoding module (new implementation).

@@ -770,8 +767,9 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
        pe = torch.cat([pe_positive, pe_negative], dim=1)
        self.pe = pe.to(device=x.device, dtype=x.dtype)

-    def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \
-            -> Tuple[torch.Tensor, torch.Tensor]:
+    def forward(
+        self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Add positional encoding.

        Args:
@@ -786,10 +784,10 @@ class EspnetRelPositionalEncoding(torch.nn.Module):
        pos_emb = self.position_encoding(size=x.size(1), offset=offset)
        return self.dropout(x), self.dropout(pos_emb)

-    def position_encoding(self,
-                          offset: Union[int, torch.Tensor],
-                          size: int) -> torch.Tensor:
-        """ For getting encoding in a streaming fashion
+    def position_encoding(
+        self, offset: Union[int, torch.Tensor], size: int
+    ) -> torch.Tensor:
+        """For getting encoding in a streaming fashion

        Attention!!!!!
        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[
            :,
-            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


-
 class LinearEmbed(torch.nn.Module):
    """Linear transform the input without subsampling

@@ -822,8 +819,9 @@ class LinearEmbed(torch.nn.Module):

    """

-    def __init__(self, idim: int, odim: int, dropout_rate: float,
-                 pos_enc_class: torch.nn.Module):
+    def __init__(
+        self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module
+    ):
        """Construct an linear object."""
        super().__init__()
        self.out = torch.nn.Sequential(
@@ -831,16 +829,15 @@ class LinearEmbed(torch.nn.Module):
            torch.nn.LayerNorm(odim, eps=1e-5),
            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],
-                          size: int) -> torch.Tensor:
+    def position_encoding(
+        self, offset: Union[int, torch.Tensor], size: int
+    ) -> torch.Tensor:
        return self.pos_enc.position_encoding(offset, size)

    def forward(
-        self,
-        x: torch.Tensor,
-        offset: Union[int, torch.Tensor] = 0
+        self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Input x.

@@ -894,16 +891,14 @@ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
    """
    batch_size = lengths.size(0)
    max_len = max_len if max_len > 0 else lengths.max().item()
-    seq_range = torch.arange(0,
-                             max_len,
-                             dtype=torch.int64,
-                             device=lengths.device)
+    seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
    seq_length_expand = lengths.unsqueeze(-1)
    mask = seq_range_expand >= seq_length_expand
    return mask

-#https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
+
+# https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
 class ConformerEncoder(torch.nn.Module):
    """Conformer encoder module."""

@@ -917,14 +912,14 @@ class ConformerEncoder(torch.nn.Module):
        dropout_rate: float = 0.1,
        positional_dropout_rate: float = 0.1,
        attention_dropout_rate: float = 0.0,
-        input_layer: str = 'linear',
-        pos_enc_layer_type: str = 'rel_pos_espnet',
+        input_layer: str = "linear",
+        pos_enc_layer_type: str = "rel_pos_espnet",
        normalize_before: bool = True,
-        static_chunk_size: int = 1, # 1: causal_mask; 0: full_mask
+        static_chunk_size: int = 1,  # 1: causal_mask; 0: full_mask
        use_dynamic_chunk: bool = False,
        use_dynamic_left_chunk: bool = False,
        positionwise_conv_kernel_size: int = 1,
-        macaron_style: bool =False,
+        macaron_style: bool = False,
        selfattention_layer_type: str = "rel_selfattn",
        activation_type: str = "swish",
        use_cnn_module: bool = False,
@@ -953,8 +948,12 @@ class ConformerEncoder(torch.nn.Module):
        """
        super().__init__()
        self.output_size = output_size
-        self.embed = LinearEmbed(input_size, output_size, dropout_rate, 
-                                        EspnetRelPositionalEncoding(output_size, positional_dropout_rate))
+        self.embed = LinearEmbed(
+            input_size,
+            output_size,
+            dropout_rate,
+            EspnetRelPositionalEncoding(output_size, positional_dropout_rate),
+        )
        self.normalize_before = normalize_before
        self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
        self.gradient_checkpointing = gradient_checkpointing
@@ -980,40 +979,60 @@ class ConformerEncoder(torch.nn.Module):
            activation,
        )
        # convolution module definition
-        convolution_layer_args = (output_size, cnn_module_kernel, activation,
-                                  cnn_module_norm, causal)
+        convolution_layer_args = (
+            output_size,
+            cnn_module_kernel,
+            activation,
+            cnn_module_norm,
+            causal,
+        )

-        self.encoders = torch.nn.ModuleList([
-            ConformerEncoderLayer(
-                output_size,
-                RelPositionMultiHeadedAttention(
-                    *encoder_selfattn_layer_args),
-                PositionwiseFeedForward(*positionwise_layer_args),
-                PositionwiseFeedForward(
-                    *positionwise_layer_args) if macaron_style else None,
-                ConvolutionModule(
-                    *convolution_layer_args) if use_cnn_module else None,
-                dropout_rate,
-                normalize_before,
-            ) for _ in range(num_blocks)
-        ])
+        self.encoders = torch.nn.ModuleList(
+            [
+                ConformerEncoderLayer(
+                    output_size,
+                    RelPositionMultiHeadedAttention(*encoder_selfattn_layer_args),
+                    PositionwiseFeedForward(*positionwise_layer_args),
+                    (
+                        PositionwiseFeedForward(*positionwise_layer_args)
+                        if macaron_style
+                        else None
+                    ),
+                    (
+                        ConvolutionModule(*convolution_layer_args)
+                        if use_cnn_module
+                        else None
+                    ),
+                    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,
-        mask_pad: torch.Tensor) -> torch.Tensor:
+        mask_pad: torch.Tensor,
+    ) -> torch.Tensor:
        for layer in self.encoders:
            xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
        return xs

    @torch.jit.unused
-    def forward_layers_checkpointed(self, xs: torch.Tensor,
-                                    chunk_masks: torch.Tensor,
-                                    pos_emb: torch.Tensor,
-                                    mask_pad: torch.Tensor) -> torch.Tensor:
+    def forward_layers_checkpointed(
+        self,
+        xs: torch.Tensor,
+        chunk_masks: torch.Tensor,
+        pos_emb: torch.Tensor,
+        mask_pad: torch.Tensor,
+    ) -> torch.Tensor:
        for layer in self.encoders:
-            xs, chunk_masks, _, _ = ckpt.checkpoint(layer.__call__, xs,
-                                                    chunk_masks, pos_emb,
-                                                    mask_pad)
+            xs, chunk_masks, _, _ = ckpt.checkpoint(
+                layer.__call__, xs, chunk_masks, pos_emb, mask_pad
+            )
        return xs

    def forward(
@@ -1050,15 +1069,17 @@ class ConformerEncoder(torch.nn.Module):
        masks = pad_mask.to(torch.bool).unsqueeze(1)  # (B, 1, T)
        xs, pos_emb = self.embed(xs)
        mask_pad = masks  # (B, 1, T/subsample_rate)
-        chunk_masks = add_optional_chunk_mask(xs, masks,
-                                              self.use_dynamic_chunk,
-                                              self.use_dynamic_left_chunk,
-                                              decoding_chunk_size,
-                                              self.static_chunk_size,
-                                              num_decoding_left_chunks) 
+        chunk_masks = add_optional_chunk_mask(
+            xs,
+            masks,
+            self.use_dynamic_chunk,
+            self.use_dynamic_left_chunk,
+            decoding_chunk_size,
+            self.static_chunk_size,
+            num_decoding_left_chunks,
+        )
        if self.gradient_checkpointing and self.training:
-            xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb,
-                                                  mask_pad)
+            xs = self.forward_layers_checkpointed(xs, chunk_masks, pos_emb, mask_pad)
        else:
            xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad)
        if self.normalize_before:
@@ -1067,4 +1088,3 @@ class ConformerEncoder(torch.nn.Module):
        # return the masks before encoder layers, and the masks will be used
        # for cross attention with decoder later
        return xs, masks
-
@@ -2,39 +2,42 @@ import re
 from opencc import OpenCC


-t2s_converter = OpenCC('t2s')
-s2t_converter = OpenCC('s2t')
+t2s_converter = OpenCC("t2s")
+s2t_converter = OpenCC("s2t")


 EMOJI_PATTERN = re.compile(
    "["
-    "\U0001F600-\U0001F64F"  # Emoticons
-    "]+", flags=re.UNICODE
+    "\U0001f600-\U0001f64f"  # Emoticons
+    "]+",
+    flags=re.UNICODE,
 )

 # 创建一个翻译表，用于替换和移除字符
-TRANSLATION_TABLE = str.maketrans({
-    '-': ' ',  # 将 '-' 替换为空格
-    ',': None,
-    '.': None,
-    '，': None,
-    '。': None,
-    '!': None,
-    '！': None,
-    '?': None,
-    '？': None,
-    '…': None,
-    ';': None,
-    '；': None,
-    ':': None,
-    '：': None,
-    '\u3000': ' ',  # 将全角空格替换为空格
-})
+TRANSLATION_TABLE = str.maketrans(
+    {
+        "-": " ",  # 将 '-' 替换为空格
+        ",": None,
+        ".": None,
+        "，": None,
+        "。": None,
+        "!": None,
+        "！": None,
+        "?": None,
+        "？": None,
+        "…": None,
+        ";": None,
+        "；": None,
+        ":": None,
+        "：": None,
+        "\u3000": " ",  # 将全角空格替换为空格
+    }
+)

 # 替换括号中的内容，包括中括号和小括号
-BACKSLASH_PATTERN = re.compile(r'\(.*?\)|\[.*?\]')
+BACKSLASH_PATTERN = re.compile(r"\(.*?\)|\[.*?\]")

-SPACE_PATTERN = re.compile('(?<!^)\s+(?!$)')
+SPACE_PATTERN = re.compile("(?<!^)\s+(?!$)")


 def normalize_text(text, language, strip=True):
@@ -45,10 +48,10 @@ def normalize_text(text, language, strip=True):
    text = text.translate(TRANSLATION_TABLE)

    # Step 2: 移除表情符号
-    text = EMOJI_PATTERN.sub('', text)
+    text = EMOJI_PATTERN.sub("", text)

    # Step 3: 连续空白字符替换为单个空格，首位除外
-    text = SPACE_PATTERN.sub(' ', text)
+    text = SPACE_PATTERN.sub(" ", text)

    # Step 4: 去除首尾空白字符（如果需要）
    if strip:
@@ -19,7 +19,7 @@ from .zh_num2words import TextNorm as zh_num2words
 from typing import Dict, List, Optional, Set, Union


-#copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py
+# copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py
 def get_spacy_lang(lang):
    if lang == "zh":
        return Chinese()
@@ -446,7 +446,9 @@ _ordinal_re = {
    "it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"),
    "pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"),
    "ar": re.compile(r"([0-9]+)(ون|ين|ث|ر|ى)"),
-    "cs": re.compile(r"([0-9]+)\.(?=\s|$)"),  # In Czech, a dot is often used after the number to indicate ordinals.
+    "cs": re.compile(
+        r"([0-9]+)\.(?=\s|$)"
+    ),  # In Czech, a dot is often used after the number to indicate ordinals.
    "ru": re.compile(r"([0-9]+)(-й|-я|-е|-ое|-ье|-го)"),
    "nl": re.compile(r"([0-9]+)(de|ste|e)"),
    "tr": re.compile(r"([0-9]+)(\.|inci|nci|uncu|üncü|\.)"),
@@ -486,7 +488,9 @@ def _expand_decimal_point(m, lang="en"):

 def _expand_currency(m, lang="en", currency="USD"):
    amount = float((re.sub(r"[^\d.]", "", m.group(0).replace(",", "."))))
-    full_amount = num2words(amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz")
+    full_amount = num2words(
+        amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz"
+    )

    and_equivalents = {
        "en": ", ",
@@ -530,13 +534,21 @@ def expand_numbers_multilingual(text, lang="en"):
        else:
            text = re.sub(_dot_number_re, _remove_dots, text)
        try:
-            text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text)
-            text = re.sub(_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text)
-            text = re.sub(_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text)
+            text = re.sub(
+                _currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text
+            )
+            text = re.sub(
+                _currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text
+            )
+            text = re.sub(
+                _currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text
+            )
        except:
            pass
        if lang != "tr":
-            text = re.sub(_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text)
+            text = re.sub(
+                _decimal_number_re, lambda m: _expand_decimal_point(m, lang), text
+            )
        text = re.sub(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
        text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
    return text
@@ -582,7 +594,15 @@ def basic_cleaners(text):

 def chinese_transliterate(text):
    return "".join(
-        [p[0] for p in pypinyin.pinyin(text, style=pypinyin.Style.TONE3, heteronym=False, neutral_tone_with_five=True)]
+        [
+            p[0]
+            for p in pypinyin.pinyin(
+                text,
+                style=pypinyin.Style.TONE3,
+                heteronym=False,
+                neutral_tone_with_five=True,
+            )
+        ]
    )


@@ -597,7 +617,9 @@ def korean_transliterate(text):
    return r.translit(text)


-DEFAULT_VOCAB_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vocab.json")
+DEFAULT_VOCAB_FILE = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "vocab.json"
+)


 class VoiceBpeTokenizer:
@@ -639,7 +661,23 @@ class VoiceBpeTokenizer:
        #     )

    def preprocess_text(self, txt, lang):
-        if lang in {"ar", "cs", "de", "en", "es", "fr", "hu", "it", "nl", "pl", "pt", "ru", "tr", "zh", "ko"}:
+        if lang in {
+            "ar",
+            "cs",
+            "de",
+            "en",
+            "es",
+            "fr",
+            "hu",
+            "it",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "zh",
+            "ko",
+        }:
            txt = multilingual_cleaners(txt, lang)
            if lang == "zh":
                txt = chinese_transliterate(txt)
@@ -672,11 +710,12 @@ class VoiceBpeTokenizer:
        # txt = txt.replace("[UNK]", "")
        return txt

-
-    #copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
+    # copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
    def batch_decode(
        self,
-        sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"],
+        sequences: Union[
+            List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"
+        ],
        skip_special_tokens: bool = False,
    ) -> List[str]:
        """
@@ -693,12 +732,9 @@ class VoiceBpeTokenizer:
        Returns:
            `List[str]`: The list of decoded sentences.
        """
-        return [
-            self.decode(seq)
-            for seq in sequences
-        ]
+        return [self.decode(seq) for seq in sequences]

-    #https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202
+    # https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202
    # def pad(self):

    def __len__(self):
@@ -716,15 +752,27 @@ def test_expand_numbers_multilingual():
        ("This is a 1st test", "This is a first test", "en"),
        ("That will be $20 sir.", "That will be twenty dollars sir.", "en"),
        ("That will be 20€ sir.", "That will be twenty euro sir.", "en"),
-        ("That will be 20.15€ sir.", "That will be twenty euro, fifteen cents sir.", "en"),
+        (
+            "That will be 20.15€ sir.",
+            "That will be twenty euro, fifteen cents sir.",
+            "en",
+        ),
        ("That's 100,000.5.", "That's one hundred thousand point five.", "en"),
        # French
        ("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"),
        ("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"),
        ("Ceci est un 1er test", "Ceci est un premier test", "fr"),
-        ("Cela vous fera $20 monsieur.", "Cela vous fera vingt dollars monsieur.", "fr"),
+        (
+            "Cela vous fera $20 monsieur.",
+            "Cela vous fera vingt dollars monsieur.",
+            "fr",
+        ),
        ("Cela vous fera 20€ monsieur.", "Cela vous fera vingt euros monsieur.", "fr"),
-        ("Cela vous fera 20,15€ monsieur.", "Cela vous fera vingt euros et quinze centimes monsieur.", "fr"),
+        (
+            "Cela vous fera 20,15€ monsieur.",
+            "Cela vous fera vingt euros et quinze centimes monsieur.",
+            "fr",
+        ),
        ("Ce sera 100.000,5.", "Ce sera cent mille virgule cinq.", "fr"),
        # German
        ("In 12,5 Sekunden.", "In zwölf Komma fünf Sekunden.", "de"),
@@ -732,21 +780,33 @@ def test_expand_numbers_multilingual():
        ("Dies ist ein 1. Test", "Dies ist ein erste Test", "de"),  # Issue with gender
        ("Das macht $20 Herr.", "Das macht zwanzig Dollar Herr.", "de"),
        ("Das macht 20€ Herr.", "Das macht zwanzig Euro Herr.", "de"),
-        ("Das macht 20,15€ Herr.", "Das macht zwanzig Euro und fünfzehn Cent Herr.", "de"),
+        (
+            "Das macht 20,15€ Herr.",
+            "Das macht zwanzig Euro und fünfzehn Cent Herr.",
+            "de",
+        ),
        # Spanish
        ("En 12,5 segundos.", "En doce punto cinco segundos.", "es"),
        ("Había 50 soldados.", "Había cincuenta soldados.", "es"),
        ("Este es un 1er test", "Este es un primero test", "es"),
        ("Eso le costará $20 señor.", "Eso le costará veinte dólares señor.", "es"),
        ("Eso le costará 20€ señor.", "Eso le costará veinte euros señor.", "es"),
-        ("Eso le costará 20,15€ señor.", "Eso le costará veinte euros con quince céntimos señor.", "es"),
+        (
+            "Eso le costará 20,15€ señor.",
+            "Eso le costará veinte euros con quince céntimos señor.",
+            "es",
+        ),
        # Italian
        ("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"),
        ("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"),
        ("Questo è un 1° test", "Questo è un primo test", "it"),
        ("Ti costerà $20 signore.", "Ti costerà venti dollari signore.", "it"),
        ("Ti costerà 20€ signore.", "Ti costerà venti euro signore.", "it"),
-        ("Ti costerà 20,15€ signore.", "Ti costerà venti euro e quindici centesimi signore.", "it"),
+        (
+            "Ti costerà 20,15€ signore.",
+            "Ti costerà venti euro e quindici centesimi signore.",
+            "it",
+        ),
        # Portuguese
        ("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
        ("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
@@ -761,8 +821,16 @@ def test_expand_numbers_multilingual():
        # Polish
        ("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"),
        ("Było 50 żołnierzy.", "Było pięćdziesiąt żołnierzy.", "pl"),
-        ("To będzie kosztować 20€ panie.", "To będzie kosztować dwadzieścia euro panie.", "pl"),
-        ("To będzie kosztować 20,15€ panie.", "To będzie kosztować dwadzieścia euro, piętnaście centów panie.", "pl"),
+        (
+            "To będzie kosztować 20€ panie.",
+            "To będzie kosztować dwadzieścia euro panie.",
+            "pl",
+        ),
+        (
+            "To będzie kosztować 20,15€ panie.",
+            "To będzie kosztować dwadzieścia euro, piętnaście centów panie.",
+            "pl",
+        ),
        # Arabic
        ("في الـ 12,5 ثانية.", "في الـ اثنا عشر  , خمسون ثانية.", "ar"),
        ("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
@@ -776,8 +844,16 @@ def test_expand_numbers_multilingual():
        # Russian
        ("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
        ("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"),
-        ("Это будет 20.15€ сэр.", "Это будет двадцать евро, пятнадцать центов сэр.", "ru"),
-        ("Это будет стоить 20€ господин.", "Это будет стоить двадцать евро господин.", "ru"),
+        (
+            "Это будет 20.15€ сэр.",
+            "Это будет двадцать евро, пятнадцать центов сэр.",
+            "ru",
+        ),
+        (
+            "Это будет стоить 20€ господин.",
+            "Это будет стоить двадцать евро господин.",
+            "ru",
+        ),
        # Dutch
        ("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"),
        ("Er waren 50 soldaten.", "Er waren vijftig soldaten.", "nl"),
@@ -817,18 +893,30 @@ def test_abbreviations_multilingual():
        ("La Dra. Martinez es muy buena.", "La doctora Martinez es muy buena.", "es"),
        # French
        ("Bonjour Mr. Dupond.", "Bonjour monsieur Dupond.", "fr"),
-        ("Mme. Moreau est absente aujourd'hui.", "madame Moreau est absente aujourd'hui.", "fr"),
+        (
+            "Mme. Moreau est absente aujourd'hui.",
+            "madame Moreau est absente aujourd'hui.",
+            "fr",
+        ),
        # German
        ("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"),
        # Portuguese
        ("Olá Sr. Silva.", "Olá senhor Silva.", "pt"),
-        ("Dra. Costa, você está disponível?", "doutora Costa, você está disponível?", "pt"),
+        (
+            "Dra. Costa, você está disponível?",
+            "doutora Costa, você está disponível?",
+            "pt",
+        ),
        # Italian
        ("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"),
        # ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern
        # Polish
        ("Dzień dobry, P. Kowalski.", "Dzień dobry, pani Kowalski.", "pl"),
-        ("M. Nowak, czy mogę zadać pytanie?", "pan Nowak, czy mogę zadać pytanie?", "pl"),
+        (
+            "M. Nowak, czy mogę zadać pytanie?",
+            "pan Nowak, czy mogę zadać pytanie?",
+            "pl",
+        ),
        # Czech
        ("P. Novák", "pan Novák", "cs"),
        ("Dr. Vojtěch", "doktor Vojtěch", "cs"),
@@ -837,7 +925,11 @@ def test_abbreviations_multilingual():
        ("Mevr. de Vries", "mevrouw de Vries", "nl"),
        # Russian
        ("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
-        ("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"),
+        (
+            "Д-р Смирнов здесь, чтобы увидеть вас.",
+            "доктор Смирнов здесь, чтобы увидеть вас.",
+            "ru",
+        ),
        # Turkish
        ("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
        ("Dr. Ayşe burada.", "doktor Ayşe burada.", "tr"),
@@ -856,8 +948,16 @@ def test_symbols_multilingual():
        ("Te veo @ la fiesta", "Te veo arroba la fiesta", "es"),
        ("J'ai 14° de fièvre", "J'ai 14 degrés de fièvre", "fr"),
        ("Die Rechnung beträgt £ 20", "Die Rechnung beträgt pfund 20", "de"),
-        ("O meu email é ana&joao@gmail.com", "O meu email é ana e joao arroba gmail.com", "pt"),
-        ("linguaggio di programmazione C#", "linguaggio di programmazione C cancelletto", "it"),
+        (
+            "O meu email é ana&joao@gmail.com",
+            "O meu email é ana e joao arroba gmail.com",
+            "pt",
+        ),
+        (
+            "linguaggio di programmazione C#",
+            "linguaggio di programmazione C cancelletto",
+            "it",
+        ),
        ("Moja temperatura to 36.6°", "Moja temperatura to 36.6 stopnie", "pl"),
        ("Mám 14% baterie", "Mám 14 procento baterie", "cs"),
        ("Těším se na tebe @ party", "Těším se na tebe na party", "cs"),
@@ -868,7 +968,11 @@ def test_symbols_multilingual():
        ("لدي 14% في البطارية", "لدي 14 في المئة في البطارية", "ar"),
        ("我的电量为 14%", "我的电量为 14 百分之", "zh"),
        ("Pilim %14 dolu.", "Pilim yüzde 14 dolu.", "tr"),
-        ("Az akkumulátorom töltöttsége 14%", "Az akkumulátorom töltöttsége 14 százalék", "hu"),
+        (
+            "Az akkumulátorom töltöttsége 14%",
+            "Az akkumulátorom töltöttsége 14 százalék",
+            "hu",
+        ),
        ("배터리 잔량이 14%입니다.", "배터리 잔량이 14 퍼센트입니다.", "ko"),
    ]

@@ -1,7 +1,7 @@
 # Authors:
 #   2019.5 Zhiyang Zhou (https://github.com/Joee1995/chn_text_norm.git)
 #   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 csv
 import os
@@ -1,3 +1,11 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
+
 import os
 import torch
 from diffusers import AutoencoderDC
@@ -21,7 +29,12 @@ VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder")

 class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
    @register_to_config
-    def __init__(self, source_sample_rate=None, dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH, vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH):
+    def __init__(
+        self,
+        source_sample_rate=None,
+        dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH,
+        vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH,
+    ):
        super(MusicDCAE, self).__init__()

        self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path)
@@ -32,9 +45,11 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):

        self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)

-        self.transform = transforms.Compose([
-            transforms.Normalize(0.5, 0.5),
-        ])
+        self.transform = transforms.Compose(
+            [
+                transforms.Normalize(0.5, 0.5),
+            ]
+        )
        self.min_mel_value = -11.0
        self.max_mel_value = 3.0
        self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
@@ -76,7 +91,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):

        max_audio_len = audio.shape[-1]
        if max_audio_len % (8 * 512) != 0:
-            audio = torch.nn.functional.pad(audio, (0, 8 * 512 - max_audio_len % (8 * 512)))
+            audio = torch.nn.functional.pad(
+                audio, (0, 8 * 512 - max_audio_len % (8 * 512))
+            )

        mels = self.forward_mel(audio)
        mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
@@ -86,7 +103,9 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
            latent = self.dcae.encoder(mel.unsqueeze(0))
            latents.append(latent)
        latents = torch.cat(latents, dim=0)
-        latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
+        latent_lengths = (
+            audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple
+        ).long()
        latents = (latents - self.shift_factor) * self.scale_factor
        return latents, latent_lengths

@@ -103,18 +122,26 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
            wav = self.vocoder.decode(mels[0]).squeeze(1)

            if sr is not None:
-                resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
+                resampler = (
+                    torchaudio.transforms.Resample(44100, sr)
+                    .to(latents.device)
+                    .to(latents.dtype)
+                )
                wav = resampler(wav)
            else:
                sr = 44100
            pred_wavs.append(wav)

        if audio_lengths is not None:
-            pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
+            pred_wavs = [
+                wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)
+            ]
        return sr, pred_wavs

    def forward(self, audios, audio_lengths=None, sr=None):
-        latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
+        latents, latent_lengths = self.encode(
+            audios=audios, audio_lengths=audio_lengths, sr=sr
+        )
        sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
        return sr, pred_wavs, latents, latent_lengths

@@ -137,5 +164,5 @@ if __name__ == "__main__":
    print("latents shape: ", latents.shape)
    print("latent_lengths: ", latent_lengths)
    print("sr: ", sr)
-    torchaudio.save("test_reconstructed.flac", pred_wavs[0], sr)
-    print("test_reconstructed.flac")
+    torchaudio.save("test_reconstructed.wav", pred_wavs[0], sr)
+    print("test_reconstructed.wav")
@@ -1,3 +1,11 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
+
 import torch
 import torch.nn as nn
 from torch import Tensor
@@ -1,3 +1,11 @@
+"""
+ACE-Step: A Step Towards Music Generation Foundation Model
+
+https://github.com/ace-step/ACE-Step
+
+Apache 2.0 License
+"""
+
 import librosa
 import torch
 from torch import nn
@@ -132,13 +140,11 @@ class ConvNeXtBlock(nn.Module):
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
        self.gamma = (
-            nn.Parameter(layer_scale_init_value *
-                         torch.ones((dim)), requires_grad=True)
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
            if layer_scale_init_value > 0
            else None
        )
-        self.drop_path = DropPath(
-            drop_path) if drop_path > 0.0 else nn.Identity()
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

    def forward(self, x, apply_residual: bool = True):
        input = x
@@ -367,8 +373,7 @@ class HiFiGANGenerator(nn.Module):
        upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
        upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
        resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
-        resblock_dilation_sizes: Tuple[Tuple[int]] = (
-            (1, 3, 5), (1, 3, 5), (1, 3, 5)),
+        resblock_dilation_sizes: Tuple[Tuple[int]] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
        num_mels: int = 128,
        upsample_initial_channel: int = 512,
        use_template: bool = True,
@@ -417,7 +422,7 @@ class HiFiGANGenerator(nn.Module):
                continue

            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(
                    Conv1d(
                        1,
@@ -498,7 +503,11 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
        upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4),
        resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13),
        resblock_dilation_sizes: Tuple[Tuple[int]] = (
-            (1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
+            (1, 3, 5),
+            (1, 3, 5),
+            (1, 3, 5),
+            (1, 3, 5),
+        ),
        num_mels: int = 512,
        upsample_initial_channel: int = 1024,
        use_template: bool = False,
@@ -565,12 +574,14 @@ class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
 if __name__ == "__main__":
    import soundfile as sf

-    x = "test_audio.flac"
-    model = ADaMoSHiFiGANV1.from_pretrained("./checkpoints/music_vocoder", local_files_only=True)
+    x = "test_audio.wav"
+    model = ADaMoSHiFiGANV1.from_pretrained(
+        "./checkpoints/music_vocoder", local_files_only=True
+    )

    wav, sr = librosa.load(x, sr=44100, mono=True)
    wav = torch.from_numpy(wav).float()[None]
    mel = model.encode(wav)

    wav = model.decode(mel)[0].mT
-    sf.write("test_audio_vocoder_rec.flac", wav.cpu().numpy(), 44100)
+    sf.write("test_audio_vocoder_rec.wav", wav.cpu().numpy(), 44100)
@@ -72,7 +72,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
        base_image_seq_len: Optional[int] = 256,
        max_image_seq_len: Optional[int] = 4096,
    ):
-        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        timesteps = np.linspace(
+            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+        )[::-1].copy()
        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)

        sigmas = timesteps / num_train_timesteps
@@ -146,7 +148,9 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):

        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
        if self.begin_index is None:
-            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps) for t in timestep
+            ]
        elif self.step_index is not None:
            # add_noise is called after first denoising step (for inpainting)
            step_indices = [self.step_index] * timestep.shape[0]
@@ -186,12 +190,16 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
        """

        if self.config.use_dynamic_shifting and mu is None:
-            raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
+            raise ValueError(
+                " you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )

        if sigmas is None:
            self.num_inference_steps = num_inference_steps
            timesteps = np.linspace(
-                self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+                self._sigma_to_t(self.sigma_max),
+                self._sigma_to_t(self.sigma_min),
+                num_inference_steps,
            )

            sigmas = timesteps / self.config.num_train_timesteps
@@ -243,7 +251,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
        s_noise: float = 1.0,
        generator: Optional[torch.Generator] = None,
        return_dict: bool = True,
-        omega: Union[float, np.array] = 0.0
+        omega: Union[float, np.array] = 0.0,
    ) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
        """
        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
@@ -67,7 +67,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
        num_train_timesteps: int = 1000,
        shift: float = 1.0,
    ):
-        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        timesteps = np.linspace(
+            1, num_train_timesteps, num_train_timesteps, dtype=np.float32
+        )[::-1].copy()
        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)

        sigmas = timesteps / num_train_timesteps
@@ -137,7 +139,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
    def _sigma_to_t(self, sigma):
        return sigma * self.config.num_train_timesteps

-    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+    def set_timesteps(
+        self, num_inference_steps: int, device: Union[str, torch.device] = None
+    ):
        """
        Sets the discrete timesteps used for the diffusion chain (to be run before inference).

@@ -150,7 +154,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
        self.num_inference_steps = num_inference_steps

        timesteps = np.linspace(
-            self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+            self._sigma_to_t(self.sigma_max),
+            self._sigma_to_t(self.sigma_min),
+            num_inference_steps,
        )

        sigmas = timesteps / self.config.num_train_timesteps
@@ -162,7 +168,9 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
        self.timesteps = timesteps.to(device=device)

        sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
-        self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
+        self.sigmas = torch.cat(
+            [sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]]
+        )

        # empty dt and derivative
        self.prev_derivative = None
@@ -208,7 +216,7 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
        s_noise: float = 1.0,
        generator: Optional[torch.Generator] = None,
        return_dict: bool = True,
-        omega: Union[float, np.array] = 0.0
+        omega: Union[float, np.array] = 0.0,
    ) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]:
        """
        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
@@ -285,13 +293,20 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
            sigma = self.sigmas[self.step_index - 1]
            sigma_next = self.sigmas[self.step_index]

-        gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0
+        gamma = (
+            min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1)
+            if s_tmin <= sigma <= s_tmax
+            else 0.0
+        )

        sigma_hat = sigma * (gamma + 1)

        if gamma > 0:
            noise = randn_tensor(
-                model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
+                model_output.shape,
+                dtype=model_output.dtype,
+                device=model_output.device,
+                generator=generator,
            )
            eps = noise * s_noise
            sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
@@ -17,6 +17,7 @@ warnings.simplefilter("ignore", category=FutureWarning)

 DEFAULT_TRAIN_PATH = "./data/example_dataset"

+
 def is_silent_audio(audio_tensor, silence_threshold=0.95):
    """
    Determine if an audio is silent and should be discarded
@@ -39,10 +40,23 @@ def is_silent_audio(audio_tensor, silence_threshold=0.95):

 # Supported languages for tokenization
 SUPPORT_LANGUAGES = {
-    "en": 259, "de": 260, "fr": 262, "es": 284, "it": 285, 
-    "pt": 286, "pl": 294, "tr": 295, "ru": 267, "cs": 293, 
-    "nl": 297, "ar": 5022, "zh": 5023, "ja": 5412, "hu": 5753,
-    "ko": 6152, "hi": 6680
+    "en": 259,
+    "de": 260,
+    "fr": 262,
+    "es": 284,
+    "it": 285,
+    "pt": 286,
+    "pl": 294,
+    "tr": 295,
+    "ru": 267,
+    "cs": 293,
+    "nl": 297,
+    "ar": 5022,
+    "zh": 5023,
+    "ja": 5412,
+    "hu": 5753,
+    "ko": 6152,
+    "hi": 6680,
 }

 # Regex pattern for structure markers like [Verse], [Chorus], etc.
@@ -54,9 +68,15 @@ class Text2MusicDataset(Dataset):
    Dataset for text-to-music generation that processes lyrics and audio files
    """

-    def __init__(self, train=True, train_dataset_path=DEFAULT_TRAIN_PATH, 
-                 max_duration=240.0, sample_size=None, shuffle=True, 
-                 minibatch_size=1):
+    def __init__(
+        self,
+        train=True,
+        train_dataset_path=DEFAULT_TRAIN_PATH,
+        max_duration=240.0,
+        sample_size=None,
+        shuffle=True,
+        minibatch_size=1,
+    ):
        """
        Initialize the Text2Music dataset

@@ -75,14 +95,107 @@ class Text2MusicDataset(Dataset):

        # Initialize language segmentation
        self.lang_segment = LangSegment()
-        self.lang_segment.setfilters([
-            'af', 'am', 'an', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'dz', 'el',
-            'en', 'eo', 'es', 'et', 'eu', 'fa', 'fi', 'fo', 'fr', 'ga', 'gl', 'gu', 'he', 'hi', 'hr', 'ht', 'hu', 'hy',
-            'id', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lb', 'lo', 'lt', 'lv', 'mg',
-            'mk', 'ml', 'mn', 'mr', 'ms', 'mt', 'nb', 'ne', 'nl', 'nn', 'no', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'qu',
-            'ro', 'ru', 'rw', 'se', 'si', 'sk', 'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'th', 'tl', 'tr', 'ug', 'uk',
-            'ur', 'vi', 'vo', 'wa', 'xh', 'zh', 'zu'
-        ])
+        self.lang_segment.setfilters(
+            [
+                "af",
+                "am",
+                "an",
+                "ar",
+                "as",
+                "az",
+                "be",
+                "bg",
+                "bn",
+                "br",
+                "bs",
+                "ca",
+                "cs",
+                "cy",
+                "da",
+                "de",
+                "dz",
+                "el",
+                "en",
+                "eo",
+                "es",
+                "et",
+                "eu",
+                "fa",
+                "fi",
+                "fo",
+                "fr",
+                "ga",
+                "gl",
+                "gu",
+                "he",
+                "hi",
+                "hr",
+                "ht",
+                "hu",
+                "hy",
+                "id",
+                "is",
+                "it",
+                "ja",
+                "jv",
+                "ka",
+                "kk",
+                "km",
+                "kn",
+                "ko",
+                "ku",
+                "ky",
+                "la",
+                "lb",
+                "lo",
+                "lt",
+                "lv",
+                "mg",
+                "mk",
+                "ml",
+                "mn",
+                "mr",
+                "ms",
+                "mt",
+                "nb",
+                "ne",
+                "nl",
+                "nn",
+                "no",
+                "oc",
+                "or",
+                "pa",
+                "pl",
+                "ps",
+                "pt",
+                "qu",
+                "ro",
+                "ru",
+                "rw",
+                "se",
+                "si",
+                "sk",
+                "sl",
+                "sq",
+                "sr",
+                "sv",
+                "sw",
+                "ta",
+                "te",
+                "th",
+                "tl",
+                "tr",
+                "ug",
+                "uk",
+                "ur",
+                "vi",
+                "vo",
+                "wa",
+                "xh",
+                "zh",
+                "zu",
+            ]
+        )

        # Initialize lyric tokenizer
        self.lyric_tokenizer = VoiceBpeTokenizer()
@@ -196,22 +309,30 @@ class Text2MusicDataset(Dataset):

                        # If debug mode, show tokenization results
                        if debug:
-                            toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
-                            logger.info(f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}")
+                            toks = self.lyric_tokenizer.batch_decode(
+                                [[tok_id] for tok_id in token_idx]
+                            )
+                            logger.info(
+                                f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}"
+                            )

                        # If tokenization contains unknown token (1), try with segment language
                        if 1 in token_idx:
                            token_idx = self.lyric_tokenizer.encode(line, lang)

                    if debug:
-                        toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
+                        toks = self.lyric_tokenizer.batch_decode(
+                            [[tok_id] for tok_id in token_idx]
+                        )
                        logger.info(f"debug {line} --> {lang} --> {toks}")

                    # Add tokens and line break
                    lyric_token_idx = lyric_token_idx + token_idx + [2]

                except Exception as e:
-                    logger.error(f"Tokenize error: {e} for line: {line}, major_language: {lang}")
+                    logger.error(
+                        f"Tokenize error: {e} for line: {line}, major_language: {lang}"
+                    )

        return lyric_token_idx

@@ -302,7 +423,9 @@ class Text2MusicDataset(Dataset):

        # Pad to minimum 3 seconds if needed
        if audio.shape[-1] < 48000 * 3:
-            audio = torch.nn.functional.pad(audio, (0, 48000 * 3 - audio.shape[-1]), 'constant', 0)
+            audio = torch.nn.functional.pad(
+                audio, (0, 48000 * 3 - audio.shape[-1]), "constant", 0
+            )

        # Check if audio is silent
        if is_silent_audio(audio):
@@ -368,9 +491,11 @@ class Text2MusicDataset(Dataset):
        lyrics = item["norm_lyrics"]
        lyrics_lines = lyrics.split("\n")
        for lyric_line in lyrics_lines:
-            candidate_lyric_chunk.append({
-                "lyric": lyric_line,
-            })
+            candidate_lyric_chunk.append(
+                {
+                    "lyric": lyric_line,
+                }
+            )

        # Limit audio length
        longest_length = 24 * 10 * 48000  # 240 seconds
@@ -482,28 +607,43 @@ class Text2MusicDataset(Dataset):
            elif k in ["src_wavs", "target_wavs", "vocal_wavs"]:
                # Pad audio to max length
                max_length = max(seq.shape[1] for seq in v)
-                padded_input_list = torch.stack([
-                    torch.nn.functional.pad(seq, (0, max_length - seq.shape[1]), 'constant', 0) 
-                    for seq in v
-                ])
+                padded_input_list = torch.stack(
+                    [
+                        torch.nn.functional.pad(
+                            seq, (0, max_length - seq.shape[1]), "constant", 0
+                        )
+                        for seq in v
+                    ]
+                )
            elif k in ["clap_conditions"]:
                # Pad time dimension of embeddings
                max_length = max(seq.shape[0] for seq in v)
                v = [
-                    torch.nn.functional.pad(seq, (0, 0, 0, max_length - seq.shape[0]), 'constant', 0) 
+                    torch.nn.functional.pad(
+                        seq, (0, 0, 0, max_length - seq.shape[0]), "constant", 0
+                    )
                    for seq in v
                ]
                padded_input_list = torch.stack(v)
            elif k == "speaker_embs":
                # Stack speaker embeddings
                padded_input_list = torch.stack(v)
-            elif k in ["chunk_masks", "clap_attention_masks", "lyric_token_ids", "lyric_masks"]:
+            elif k in [
+                "chunk_masks",
+                "clap_attention_masks",
+                "lyric_token_ids",
+                "lyric_masks",
+            ]:
                # Pad sequence tensors
                max_length = max(len(seq) for seq in v)
-                padded_input_list = torch.stack([
-                    torch.nn.functional.pad(seq, (0, max_length - len(seq)), 'constant', 0) 
-                    for seq in v
-                ])
+                padded_input_list = torch.stack(
+                    [
+                        torch.nn.functional.pad(
+                            seq, (0, max_length - len(seq)), "constant", 0
+                        )
+                        for seq in v
+                    ]
+                )

            output[k] = padded_input_list

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

@@ -41,8 +49,6 @@ In this moment we take flight
 """


-
-
 def create_output_ui(task_name="Text2Music"):
    # For many consumer-grade GPU devices, only one batch can be run
    output_audio1 = gr.Audio(type="filepath", label=f"{task_name} Generated Audio 1")
@@ -68,41 +74,162 @@ def create_text2music_ui(
        with gr.Column():
            with gr.Row(equal_height=True):
                # add markdown, tags and lyrics examples are from ai music generation community
-                audio_duration = gr.Slider(-1, 240.0, step=0.00001, value=-1, label="Audio Duration", interactive=True, info="-1 means random duration (30 ~ 240).", scale=9)
+                audio_duration = gr.Slider(
+                    -1,
+                    240.0,
+                    step=0.00001,
+                    value=-1,
+                    label="Audio Duration",
+                    interactive=True,
+                    info="-1 means random duration (30 ~ 240).",
+                    scale=9,
+                )
                sample_bnt = gr.Button("Sample", variant="primary", scale=1)

-            prompt = gr.Textbox(lines=2, label="Tags", max_lines=4, value=TAG_DEFAULT, info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community")
-            lyrics = gr.Textbox(lines=9, label="Lyrics", max_lines=13, value=LYRIC_DEFAULT, info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics")
+            prompt = gr.Textbox(
+                lines=2,
+                label="Tags",
+                max_lines=4,
+                value=TAG_DEFAULT,
+                info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community",
+            )
+            lyrics = gr.Textbox(
+                lines=9,
+                label="Lyrics",
+                max_lines=13,
+                value=LYRIC_DEFAULT,
+                info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics",
+            )

            with gr.Accordion("Basic Settings", open=False):
-                infer_step = gr.Slider(minimum=1, maximum=1000, step=1, value=27, label="Infer Steps", interactive=True)
-                guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=15.0, label="Guidance Scale", interactive=True, info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.")
-                guidance_scale_text = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Text", interactive=True, info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start")
-                guidance_scale_lyric = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Lyric", interactive=True)
+                infer_step = gr.Slider(
+                    minimum=1,
+                    maximum=1000,
+                    step=1,
+                    value=27,
+                    label="Infer Steps",
+                    interactive=True,
+                )
+                guidance_scale = gr.Slider(
+                    minimum=0.0,
+                    maximum=200.0,
+                    step=0.1,
+                    value=15.0,
+                    label="Guidance Scale",
+                    interactive=True,
+                    info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.",
+                )
+                guidance_scale_text = gr.Slider(
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                    label="Guidance Scale Text",
+                    interactive=True,
+                    info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start",
+                )
+                guidance_scale_lyric = gr.Slider(
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                    label="Guidance Scale Lyric",
+                    interactive=True,
+                )

-                manual_seeds = gr.Textbox(label="manual seeds (default None)", placeholder="1,2,3,4", value=None, info="Seed for the generation")
+                manual_seeds = gr.Textbox(
+                    label="manual seeds (default None)",
+                    placeholder="1,2,3,4",
+                    value=None,
+                    info="Seed for the generation",
+                )

            with gr.Accordion("Advanced Settings", open=False):
-                scheduler_type = gr.Radio(["euler", "heun"], value="euler", label="Scheduler Type", elem_id="scheduler_type", info="Scheduler type for the generation. euler is recommended. heun will take more time.")
-                cfg_type = gr.Radio(["cfg", "apg", "cfg_star"], value="apg", label="CFG Type", elem_id="cfg_type", info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.")
-                use_erg_tag = gr.Checkbox(label="use ERG for tag", value=True, info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.")
-                use_erg_lyric = gr.Checkbox(label="use ERG for lyric", value=True, info="The same but apply to lyric encoder's attention.")
-                use_erg_diffusion = gr.Checkbox(label="use ERG for diffusion", value=True, info="The same but apply to diffusion model's attention.")
+                scheduler_type = gr.Radio(
+                    ["euler", "heun"],
+                    value="euler",
+                    label="Scheduler Type",
+                    elem_id="scheduler_type",
+                    info="Scheduler type for the generation. euler is recommended. heun will take more time.",
+                )
+                cfg_type = gr.Radio(
+                    ["cfg", "apg", "cfg_star"],
+                    value="apg",
+                    label="CFG Type",
+                    elem_id="cfg_type",
+                    info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.",
+                )
+                use_erg_tag = gr.Checkbox(
+                    label="use ERG for tag",
+                    value=True,
+                    info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.",
+                )
+                use_erg_lyric = gr.Checkbox(
+                    label="use ERG for lyric",
+                    value=True,
+                    info="The same but apply to lyric encoder's attention.",
+                )
+                use_erg_diffusion = gr.Checkbox(
+                    label="use ERG for diffusion",
+                    value=True,
+                    info="The same but apply to diffusion model's attention.",
+                )

-                omega_scale = gr.Slider(minimum=-100.0, maximum=100.0, step=0.1, value=10.0, label="Granularity Scale", interactive=True, info="Granularity scale for the generation. Higher values can reduce artifacts")
+                omega_scale = gr.Slider(
+                    minimum=-100.0,
+                    maximum=100.0,
+                    step=0.1,
+                    value=10.0,
+                    label="Granularity Scale",
+                    interactive=True,
+                    info="Granularity scale for the generation. Higher values can reduce artifacts",
+                )

-                guidance_interval = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Guidance Interval", interactive=True, info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)")
-                guidance_interval_decay = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Guidance Interval Decay", interactive=True, info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.")
-                min_guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=3.0, label="Min Guidance Scale", interactive=True, info="Min guidance scale for guidance interval decay's end scale")
-                oss_steps = gr.Textbox(label="OSS Steps", placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200", value=None, info="Optimal Steps for the generation. But not test well")
+                guidance_interval = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.5,
+                    label="Guidance Interval",
+                    interactive=True,
+                    info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)",
+                )
+                guidance_interval_decay = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.0,
+                    label="Guidance Interval Decay",
+                    interactive=True,
+                    info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.",
+                )
+                min_guidance_scale = gr.Slider(
+                    minimum=0.0,
+                    maximum=200.0,
+                    step=0.1,
+                    value=3.0,
+                    label="Min Guidance Scale",
+                    interactive=True,
+                    info="Min guidance scale for guidance interval decay's end scale",
+                )
+                oss_steps = gr.Textbox(
+                    label="OSS Steps",
+                    placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200",
+                    value=None,
+                    info="Optimal Steps for the generation. But not test well",
+                )

            text2music_bnt = gr.Button("Generate", variant="primary")

        with gr.Column():
            outputs, input_params_json = create_output_ui()
            with gr.Tab("retake"):
-                retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
-                retake_seeds = gr.Textbox(label="retake seeds (default None)", placeholder="", value=None)
+                retake_variance = gr.Slider(
+                    minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
+                )
+                retake_seeds = gr.Textbox(
+                    label="retake seeds (default None)", placeholder="", value=None
+                )
                retake_bnt = gr.Button("Retake", variant="primary")
                retake_outputs, retake_input_params_json = create_output_ui("Retake")

@@ -124,8 +251,16 @@ def create_text2music_ui(
                        json_data["use_erg_lyric"],
                        json_data["use_erg_diffusion"],
                        ", ".join(map(str, json_data["oss_steps"])),
-                        json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
-                        json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
+                        (
+                            json_data["guidance_scale_text"]
+                            if "guidance_scale_text" in json_data
+                            else 0.0
+                        ),
+                        (
+                            json_data["guidance_scale_lyric"]
+                            if "guidance_scale_lyric" in json_data
+                            else 0.0
+                        ),
                        retake_seeds=retake_seeds,
                        retake_variance=retake_variance,
                        task="retake",
@@ -141,15 +276,45 @@ def create_text2music_ui(
                    outputs=retake_outputs + [retake_input_params_json],
                )
            with gr.Tab("repainting"):
-                retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
-                retake_seeds = gr.Textbox(label="repaint seeds (default None)", placeholder="", value=None)
-                repaint_start = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Repaint Start Time", interactive=True)
-                repaint_end = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Repaint End Time", interactive=True)
-                repaint_source = gr.Radio(["text2music", "last_repaint", "upload"], value="text2music", label="Repaint Source", elem_id="repaint_source")
+                retake_variance = gr.Slider(
+                    minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance"
+                )
+                retake_seeds = gr.Textbox(
+                    label="repaint seeds (default None)", placeholder="", value=None
+                )
+                repaint_start = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=0.0,
+                    label="Repaint Start Time",
+                    interactive=True,
+                )
+                repaint_end = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=30.0,
+                    label="Repaint End Time",
+                    interactive=True,
+                )
+                repaint_source = gr.Radio(
+                    ["text2music", "last_repaint", "upload"],
+                    value="text2music",
+                    label="Repaint Source",
+                    elem_id="repaint_source",
+                )

-                repaint_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="repaint_source_audio_upload")
+                repaint_source_audio_upload = gr.Audio(
+                    label="Upload Audio",
+                    type="filepath",
+                    visible=False,
+                    elem_id="repaint_source_audio_upload",
+                )
                repaint_source.change(
-                    fn=lambda x: gr.update(visible=x == "upload", elem_id="repaint_source_audio_upload"),
+                    fn=lambda x: gr.update(
+                        visible=x == "upload", elem_id="repaint_source_audio_upload"
+                    ),
                    inputs=[repaint_source],
                    outputs=[repaint_source_audio_upload],
                )
@@ -187,9 +352,7 @@ def create_text2music_ui(
                    if repaint_source == "upload":
                        src_audio_path = repaint_source_audio_upload
                        audio_duration = librosa.get_duration(filename=src_audio_path)
-                        json_data = {
-                            "audio_duration": audio_duration
-                        }
+                        json_data = {"audio_duration": audio_duration}
                    elif repaint_source == "text2music":
                        json_data = text2music_json_data
                        src_audio_path = json_data["audio_path"]
@@ -258,11 +421,33 @@ def create_text2music_ui(
            with gr.Tab("edit"):
                edit_prompt = gr.Textbox(lines=2, label="Edit Tags", max_lines=4)
                edit_lyrics = gr.Textbox(lines=9, label="Edit Lyrics", max_lines=13)
-                retake_seeds = gr.Textbox(label="edit seeds (default None)", placeholder="", value=None)
+                retake_seeds = gr.Textbox(
+                    label="edit seeds (default None)", placeholder="", value=None
+                )

-                edit_type = gr.Radio(["only_lyrics", "remix"], value="only_lyrics", label="Edit Type", elem_id="edit_type", info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre")
-                edit_n_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.6, label="edit_n_min", interactive=True)
-                edit_n_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="edit_n_max", interactive=True)
+                edit_type = gr.Radio(
+                    ["only_lyrics", "remix"],
+                    value="only_lyrics",
+                    label="Edit Type",
+                    elem_id="edit_type",
+                    info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre",
+                )
+                edit_n_min = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.6,
+                    label="edit_n_min",
+                    interactive=True,
+                )
+                edit_n_max = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=1.0,
+                    label="edit_n_max",
+                    interactive=True,
+                )

                def edit_type_change_func(edit_type):
                    if edit_type == "only_lyrics":
@@ -276,13 +461,25 @@ def create_text2music_ui(
                edit_type.change(
                    edit_type_change_func,
                    inputs=[edit_type],
-                    outputs=[edit_n_min, edit_n_max]
+                    outputs=[edit_n_min, edit_n_max],
                )

-                edit_source = gr.Radio(["text2music", "last_edit", "upload"], value="text2music", label="Edit Source", elem_id="edit_source")
-                edit_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="edit_source_audio_upload")
+                edit_source = gr.Radio(
+                    ["text2music", "last_edit", "upload"],
+                    value="text2music",
+                    label="Edit Source",
+                    elem_id="edit_source",
+                )
+                edit_source_audio_upload = gr.Audio(
+                    label="Upload Audio",
+                    type="filepath",
+                    visible=False,
+                    elem_id="edit_source_audio_upload",
+                )
                edit_source.change(
-                    fn=lambda x: gr.update(visible=x == "upload", elem_id="edit_source_audio_upload"),
+                    fn=lambda x: gr.update(
+                        visible=x == "upload", elem_id="edit_source_audio_upload"
+                    ),
                    inputs=[edit_source],
                    outputs=[edit_source_audio_upload],
                )
@@ -321,9 +518,7 @@ def create_text2music_ui(
                    if edit_source == "upload":
                        src_audio_path = edit_source_audio_upload
                        audio_duration = librosa.get_duration(filename=src_audio_path)
-                        json_data = {
-                            "audio_duration": audio_duration
-                        }
+                        json_data = {"audio_duration": audio_duration}
                    elif edit_source == "text2music":
                        json_data = text2music_json_data
                        src_audio_path = json_data["audio_path"]
@@ -397,14 +592,42 @@ def create_text2music_ui(
                    outputs=edit_outputs + [edit_input_params_json],
                )
            with gr.Tab("extend"):
-                extend_seeds = gr.Textbox(label="extend seeds (default None)", placeholder="", value=None)
-                left_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Left Extend Length", interactive=True)
-                right_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Right Extend Length", interactive=True)
-                extend_source = gr.Radio(["text2music", "last_extend", "upload"], value="text2music", label="Extend Source", elem_id="extend_source")
+                extend_seeds = gr.Textbox(
+                    label="extend seeds (default None)", placeholder="", value=None
+                )
+                left_extend_length = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=0.0,
+                    label="Left Extend Length",
+                    interactive=True,
+                )
+                right_extend_length = gr.Slider(
+                    minimum=0.0,
+                    maximum=240.0,
+                    step=0.01,
+                    value=30.0,
+                    label="Right Extend Length",
+                    interactive=True,
+                )
+                extend_source = gr.Radio(
+                    ["text2music", "last_extend", "upload"],
+                    value="text2music",
+                    label="Extend Source",
+                    elem_id="extend_source",
+                )

-                extend_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="extend_source_audio_upload")
+                extend_source_audio_upload = gr.Audio(
+                    label="Upload Audio",
+                    type="filepath",
+                    visible=False,
+                    elem_id="extend_source_audio_upload",
+                )
                extend_source.change(
-                    fn=lambda x: gr.update(visible=x == "upload", elem_id="extend_source_audio_upload"),
+                    fn=lambda x: gr.update(
+                        visible=x == "upload", elem_id="extend_source_audio_upload"
+                    ),
                    inputs=[extend_source],
                    outputs=[extend_source_audio_upload],
                )
@@ -442,9 +665,7 @@ def create_text2music_ui(
                        src_audio_path = extend_source_audio_upload
                        # get audio duration
                        audio_duration = librosa.get_duration(filename=src_audio_path)
-                        json_data = {
-                            "audio_duration": audio_duration
-                        }
+                        json_data = {"audio_duration": audio_duration}
                    elif extend_source == "text2music":
                        json_data = text2music_json_data
                        src_audio_path = json_data["audio_path"]
@@ -531,8 +752,16 @@ def create_text2music_ui(
                json_data["use_erg_lyric"],
                json_data["use_erg_diffusion"],
                ", ".join(map(str, json_data["oss_steps"])),
-                json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
-                json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
+                (
+                    json_data["guidance_scale_text"]
+                    if "guidance_scale_text" in json_data
+                    else 0.0
+                ),
+                (
+                    json_data["guidance_scale_lyric"]
+                    if "guidance_scale_lyric" in json_data
+                    else 0.0
+                ),
            )

        sample_bnt.click(
@@ -580,7 +809,8 @@ def create_text2music_ui(
            oss_steps,
            guidance_scale_text,
            guidance_scale_lyric,
-        ], outputs=outputs + [input_params_json]
+        ],
+        outputs=outputs + [input_params_json],
    )


@@ -594,7 +824,8 @@ def create_main_demo_ui(
        gr.Markdown(
            """
            <h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
-        """)
+        """
+        )

        with gr.Tab("text2music"):
            create_text2music_ui(
@@ -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)
@@ -1,48 +1,54 @@
-import argparse
-parser = argparse.ArgumentParser()
-parser.add_argument("--checkpoint_path", type=str, default="")
-parser.add_argument("--bf16", type=bool, default=True)
-parser.add_argument("--torch_compile", type=bool, default=False)
-parser.add_argument("--device_id", type=int, default=0)
-parser.add_argument("--output_path", type=str, default=None)
-args = parser.parse_args()
-
+import click
 import os

-os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
-
-from pipeline_ace_step import ACEStepPipeline
-from data_sampler import DataSampler
+from acestep.pipeline_ace_step import ACEStepPipeline
+from acestep.data_sampler import DataSampler


 def sample_data(json_data):
    return (
-            json_data["audio_duration"],
-            json_data["prompt"],
-            json_data["lyrics"],
-            json_data["infer_step"],
-            json_data["guidance_scale"],
-            json_data["scheduler_type"],
-            json_data["cfg_type"],
-            json_data["omega_scale"],
-            ", ".join(map(str, json_data["actual_seeds"])),
-            json_data["guidance_interval"],
-            json_data["guidance_interval_decay"],
-            json_data["min_guidance_scale"],
-            json_data["use_erg_tag"],
-            json_data["use_erg_lyric"],
-            json_data["use_erg_diffusion"],
-            ", ".join(map(str, json_data["oss_steps"])),
-            json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
-            json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
-            )
+        json_data["audio_duration"],
+        json_data["prompt"],
+        json_data["lyrics"],
+        json_data["infer_step"],
+        json_data["guidance_scale"],
+        json_data["scheduler_type"],
+        json_data["cfg_type"],
+        json_data["omega_scale"],
+        ", ".join(map(str, json_data["actual_seeds"])),
+        json_data["guidance_interval"],
+        json_data["guidance_interval_decay"],
+        json_data["min_guidance_scale"],
+        json_data["use_erg_tag"],
+        json_data["use_erg_lyric"],
+        json_data["use_erg_diffusion"],
+        ", ".join(map(str, json_data["oss_steps"])),
+        json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
+        (
+            json_data["guidance_scale_lyric"]
+            if "guidance_scale_lyric" in json_data
+            else 0.0
+        ),
+    )

-def main(args):
+
+@click.command()
+@click.option(
+    "--checkpoint_path", type=str, default="", help="Path to the checkpoint directory"
+)
+@click.option("--bf16", type=bool, default=True, help="Whether to use bfloat16")
+@click.option(
+    "--torch_compile", type=bool, default=False, help="Whether to use torch compile"
+)
+@click.option("--device_id", type=int, default=0, help="Device ID to use")
+@click.option("--output_path", type=str, default=None, help="Path to save the output")
+def main(checkpoint_path, bf16, torch_compile, device_id, output_path):
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)

    model_demo = ACEStepPipeline(
-        checkpoint_dir=args.checkpoint_path,
-        dtype="bfloat16" if args.bf16 else "float32",
-        torch_compile=args.torch_compile
+        checkpoint_dir=checkpoint_path,
+        dtype="bfloat16" if bf16 else "float32",
+        torch_compile=torch_compile,
    )
    print(model_demo)

@@ -52,46 +58,49 @@ def main(args):
    json_data = sample_data(json_data)
    print(json_data)

+    (
+        audio_duration,
+        prompt,
+        lyrics,
+        infer_step,
+        guidance_scale,
+        scheduler_type,
+        cfg_type,
+        omega_scale,
+        manual_seeds,
+        guidance_interval,
+        guidance_interval_decay,
+        min_guidance_scale,
+        use_erg_tag,
+        use_erg_lyric,
+        use_erg_diffusion,
+        oss_steps,
+        guidance_scale_text,
+        guidance_scale_lyric,
+    ) = json_data

-    audio_duration,\
-    prompt, \
-    lyrics,\
-    infer_step, \
-    guidance_scale,\
-    scheduler_type, \
-    cfg_type, \
-    omega_scale, \
-    manual_seeds, \
-    guidance_interval, \
-    guidance_interval_decay, \
-    min_guidance_scale, \
-    use_erg_tag, \
-    use_erg_lyric, \
-    use_erg_diffusion, \
-    oss_steps, \
-    guidance_scale_text, \
-    guidance_scale_lyric = json_data
+    model_demo(
+        audio_duration,
+        prompt,
+        lyrics,
+        infer_step,
+        guidance_scale,
+        scheduler_type,
+        cfg_type,
+        omega_scale,
+        manual_seeds,
+        guidance_interval,
+        guidance_interval_decay,
+        min_guidance_scale,
+        use_erg_tag,
+        use_erg_lyric,
+        use_erg_diffusion,
+        oss_steps,
+        guidance_scale_text,
+        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__":
-    main(args)
+    main()
@@ -1,9 +0,0 @@
-from .LangSegment import LangSegment
-
-
-# release
-__version__ = '0.3.5'
-
-
-# develop
-__develop__ = 'dev-0.0.1'
@@ -20,3 +20,4 @@ spacy==3.8.4
 accelerate==1.6.0
 cutlet
 fugashi[unidic-lite]
+click
@@ -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",
+        ],
+    },
+)
@@ -10,23 +10,27 @@ import torch.nn.functional as F
 import torch.utils.data
 from pytorch_lightning.core import LightningModule
 from torch.utils.data import DataLoader
-from schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from text2music_dataset import Text2MusicDataset
+from acestep.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from acestep.text2music_dataset import Text2MusicDataset
 from loguru import logger
 from transformers import AutoModel, Wav2Vec2FeatureExtractor
 import torchaudio
-from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
+from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import (
+    retrieve_timesteps,
+)
 from diffusers.utils.torch_utils import randn_tensor
-from apg_guidance import apg_forward, MomentumBuffer
+from acestep.apg_guidance import apg_forward, MomentumBuffer
 from tqdm import tqdm
 import random
 import os
-from pipeline_ace_step import ACEStepPipeline
+from acestep.pipeline_ace_step import ACEStepPipeline


 matplotlib.use("Agg")
 torch.backends.cudnn.benchmark = False
-torch.set_float32_matmul_precision('high')
+torch.set_float32_matmul_precision("high")


 class Pipeline(LightningModule):
@@ -47,7 +51,7 @@ class Pipeline(LightningModule):
        max_steps: int = 200000,
        warmup_steps: int = 4000,
        dataset_path: str = "./data/your_dataset_path",
-        lora_config_path: str = None
+        lora_config_path: str = None,
    ):
        super().__init__()

@@ -91,15 +95,31 @@ class Pipeline(LightningModule):
        if self.is_train:
            self.transformers.train()

-            self.mert_model = AutoModel.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True, cache_dir=checkpoint_dir).eval()
+            self.mert_model = AutoModel.from_pretrained(
+                "m-a-p/MERT-v1-330M", trust_remote_code=True, cache_dir=checkpoint_dir
+            ).eval()
            self.mert_model.requires_grad_(False)
-            self.resampler_mert = torchaudio.transforms.Resample(orig_freq=48000, new_freq=24000)
-            self.processor_mert = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True)
+            self.resampler_mert = torchaudio.transforms.Resample(
+                orig_freq=48000, new_freq=24000
+            )
+            self.processor_mert = Wav2Vec2FeatureExtractor.from_pretrained(
+                "m-a-p/MERT-v1-330M", trust_remote_code=True
+            )

-            self.hubert_model = AutoModel.from_pretrained("utter-project/mHuBERT-147", local_files_only=True, cache_dir=checkpoint_dir).eval()
+            self.hubert_model = AutoModel.from_pretrained(
+                "utter-project/mHuBERT-147",
+                local_files_only=True,
+                cache_dir=checkpoint_dir,
+            ).eval()
            self.hubert_model.requires_grad_(False)
-            self.resampler_mhubert = torchaudio.transforms.Resample(orig_freq=48000, new_freq=16000)
-            self.processor_mhubert = Wav2Vec2FeatureExtractor.from_pretrained("utter-project/mHuBERT-147", local_files_only=True, cache_dir=checkpoint_dir)
+            self.resampler_mhubert = torchaudio.transforms.Resample(
+                orig_freq=48000, new_freq=16000
+            )
+            self.processor_mhubert = Wav2Vec2FeatureExtractor.from_pretrained(
+                "utter-project/mHuBERT-147",
+                local_files_only=True,
+                cache_dir=checkpoint_dir,
+            )

            self.ssl_coeff = ssl_coeff

@@ -110,9 +130,21 @@ class Pipeline(LightningModule):
        actual_lengths_24k = wav_lengths // 2  # 48kHz -> 24kHz

        # Normalize the actual audio part
-        means = torch.stack([mert_input_wavs_mono_24k[i, :actual_lengths_24k[i]].mean() for i in range(bsz)])
-        vars = torch.stack([mert_input_wavs_mono_24k[i, :actual_lengths_24k[i]].var() for i in range(bsz)])
-        mert_input_wavs_mono_24k = (mert_input_wavs_mono_24k - means.view(-1, 1)) / torch.sqrt(vars.view(-1, 1) + 1e-7)
+        means = torch.stack(
+            [
+                mert_input_wavs_mono_24k[i, : actual_lengths_24k[i]].mean()
+                for i in range(bsz)
+            ]
+        )
+        vars = torch.stack(
+            [
+                mert_input_wavs_mono_24k[i, : actual_lengths_24k[i]].var()
+                for i in range(bsz)
+            ]
+        )
+        mert_input_wavs_mono_24k = (
+            mert_input_wavs_mono_24k - means.view(-1, 1)
+        ) / torch.sqrt(vars.view(-1, 1) + 1e-7)

        # MERT SSL constraint
        # Define the length of each chunk (5 seconds of samples)
@@ -131,7 +163,9 @@ class Pipeline(LightningModule):
                end = min(start + chunk_size, actual_length)
                chunk = audio[start:end]
                if len(chunk) < chunk_size:
-                    chunk = F.pad(chunk, (0, chunk_size - len(chunk)))  # Pad insufficient parts with zeros
+                    chunk = F.pad(
+                        chunk, (0, chunk_size - len(chunk))
+                    )  # Pad insufficient parts with zeros
                all_chunks.append(chunk)
                chunk_actual_lengths.append(end - start)

@@ -147,14 +181,21 @@ class Pipeline(LightningModule):
        chunk_num_features = [(length + 319) // 320 for length in chunk_actual_lengths]

        # Trim the hidden states of each chunk
-        chunk_hidden_states = [mert_ssl_hidden_states[i, :chunk_num_features[i], :] for i in range(len(all_chunks))]
+        chunk_hidden_states = [
+            mert_ssl_hidden_states[i, : chunk_num_features[i], :]
+            for i in range(len(all_chunks))
+        ]

        # Organize hidden states by audio
        mert_ssl_hidden_states_list = []
        chunk_idx = 0
        for i in range(bsz):
-            audio_chunks = chunk_hidden_states[chunk_idx:chunk_idx + num_chunks_per_audio[i]]
-            audio_hidden = torch.cat(audio_chunks, dim=0)  # Concatenate chunks of the same audio
+            audio_chunks = chunk_hidden_states[
+                chunk_idx : chunk_idx + num_chunks_per_audio[i]
+            ]
+            audio_hidden = torch.cat(
+                audio_chunks, dim=0
+            )  # Concatenate chunks of the same audio
            mert_ssl_hidden_states_list.append(audio_hidden)
            chunk_idx += num_chunks_per_audio[i]

@@ -168,18 +209,29 @@ class Pipeline(LightningModule):
        actual_lengths_16k = wav_lengths // 3  # Convert lengths from 48kHz to 16kHz

        # Step 2: Zero-mean unit-variance normalization (only on actual audio)
-        means = torch.stack([mhubert_input_wavs_mono_16k[i, :actual_lengths_16k[i]].mean() 
-                            for i in range(bsz)])
-        vars = torch.stack([mhubert_input_wavs_mono_16k[i, :actual_lengths_16k[i]].var() 
-                            for i in range(bsz)])
-        mhubert_input_wavs_mono_16k = (mhubert_input_wavs_mono_16k - means.view(-1, 1)) / \
-                                    torch.sqrt(vars.view(-1, 1) + 1e-7)
+        means = torch.stack(
+            [
+                mhubert_input_wavs_mono_16k[i, : actual_lengths_16k[i]].mean()
+                for i in range(bsz)
+            ]
+        )
+        vars = torch.stack(
+            [
+                mhubert_input_wavs_mono_16k[i, : actual_lengths_16k[i]].var()
+                for i in range(bsz)
+            ]
+        )
+        mhubert_input_wavs_mono_16k = (
+            mhubert_input_wavs_mono_16k - means.view(-1, 1)
+        ) / torch.sqrt(vars.view(-1, 1) + 1e-7)

        # Step 3: Define chunk size for MHubert (30 seconds at 16kHz)
        chunk_size = 16000 * 30  # 30 seconds = 480,000 samples

        # Step 4: Split audio into chunks
-        num_chunks_per_audio = (actual_lengths_16k + chunk_size - 1) // chunk_size  # Ceiling division
+        num_chunks_per_audio = (
+            actual_lengths_16k + chunk_size - 1
+        ) // chunk_size  # Ceiling division
        all_chunks = []
        chunk_actual_lengths = []

@@ -206,20 +258,33 @@ class Pipeline(LightningModule):
        chunk_num_features = [(length + 319) // 320 for length in chunk_actual_lengths]

        # Step 8: Trim hidden states to remove padding effects
-        chunk_hidden_states = [mhubert_ssl_hidden_states[i, :chunk_num_features[i], :] for i in range(len(all_chunks))]
+        chunk_hidden_states = [
+            mhubert_ssl_hidden_states[i, : chunk_num_features[i], :]
+            for i in range(len(all_chunks))
+        ]

        # Step 9: Reorganize hidden states by original audio
        mhubert_ssl_hidden_states_list = []
        chunk_idx = 0
        for i in range(bsz):
-            audio_chunks = chunk_hidden_states[chunk_idx:chunk_idx + num_chunks_per_audio[i]]
-            audio_hidden = torch.cat(audio_chunks, dim=0)  # Concatenate chunks for this audio
+            audio_chunks = chunk_hidden_states[
+                chunk_idx : chunk_idx + num_chunks_per_audio[i]
+            ]
+            audio_hidden = torch.cat(
+                audio_chunks, dim=0
+            )  # Concatenate chunks for this audio
            mhubert_ssl_hidden_states_list.append(audio_hidden)
            chunk_idx += num_chunks_per_audio[i]
        return mhubert_ssl_hidden_states_list

    def get_text_embeddings(self, texts, device, text_max_length=256):
-        inputs = self.text_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=text_max_length)
+        inputs = self.text_tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=text_max_length,
+        )
        inputs = {key: value.to(device) for key, value in inputs.items()}
        if self.text_encoder_model.device != device:
            self.text_encoder_model.to(device)
@@ -243,15 +308,21 @@ class Pipeline(LightningModule):
        if train:
            with torch.amp.autocast(device_type="cuda", dtype=dtype):
                mert_ssl_hidden_states = self.infer_mert_ssl(target_wavs, wav_lengths)
-                mhubert_ssl_hidden_states = self.infer_mhubert_ssl(target_wavs, wav_lengths)
+                mhubert_ssl_hidden_states = self.infer_mhubert_ssl(
+                    target_wavs, wav_lengths
+                )

        # 1: text embedding
        texts = batch["prompts"]
-        encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(texts, device)
+        encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(
+            texts, device
+        )
        encoder_text_hidden_states = encoder_text_hidden_states.to(dtype)

        target_latents, _ = self.dcae.encode(target_wavs, wav_lengths)
-        attention_mask = torch.ones(bs, target_latents.shape[-1], device=device, dtype=dtype)
+        attention_mask = torch.ones(
+            bs, target_latents.shape[-1], device=device, dtype=dtype
+        )

        speaker_embds = batch["speaker_embs"].to(dtype)
        keys = batch["keys"]
@@ -263,27 +334,43 @@ class Pipeline(LightningModule):
            full_cfg_condition_mask = torch.where(
                (torch.rand(size=(bs,), device=device) < 0.15),
                torch.zeros(size=(bs,), device=device),
-                torch.ones(size=(bs,), device=device)
+                torch.ones(size=(bs,), device=device),
            ).long()
            # N x T x 768
-            encoder_text_hidden_states = torch.where(full_cfg_condition_mask.unsqueeze(1).unsqueeze(1).bool(), encoder_text_hidden_states, torch.zeros_like(encoder_text_hidden_states))
+            encoder_text_hidden_states = torch.where(
+                full_cfg_condition_mask.unsqueeze(1).unsqueeze(1).bool(),
+                encoder_text_hidden_states,
+                torch.zeros_like(encoder_text_hidden_states),
+            )

            full_cfg_condition_mask = torch.where(
                (torch.rand(size=(bs,), device=device) < 0.50),
                torch.zeros(size=(bs,), device=device),
-                torch.ones(size=(bs,), device=device)
+                torch.ones(size=(bs,), device=device),
            ).long()
            # N x 512
-            speaker_embds = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), speaker_embds, torch.zeros_like(speaker_embds))
+            speaker_embds = torch.where(
+                full_cfg_condition_mask.unsqueeze(1).bool(),
+                speaker_embds,
+                torch.zeros_like(speaker_embds),
+            )

            # Lyrics
            full_cfg_condition_mask = torch.where(
                (torch.rand(size=(bs,), device=device) < 0.15),
                torch.zeros(size=(bs,), device=device),
-                torch.ones(size=(bs,), device=device)
+                torch.ones(size=(bs,), device=device),
            ).long()
-            lyric_token_ids = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), lyric_token_ids, torch.zeros_like(lyric_token_ids))
-            lyric_mask = torch.where(full_cfg_condition_mask.unsqueeze(1).bool(), lyric_mask, torch.zeros_like(lyric_mask))
+            lyric_token_ids = torch.where(
+                full_cfg_condition_mask.unsqueeze(1).bool(),
+                lyric_token_ids,
+                torch.zeros_like(lyric_token_ids),
+            )
+            lyric_mask = torch.where(
+                full_cfg_condition_mask.unsqueeze(1).bool(),
+                lyric_mask,
+                torch.zeros_like(lyric_mask),
+            )

        return (
            keys,
@@ -305,10 +392,12 @@ class Pipeline(LightningModule):
        )

    def configure_optimizers(self):
-        trainable_params = [p for name, p in self.transformers.named_parameters() if p.requires_grad]
+        trainable_params = [
+            p for name, p in self.transformers.named_parameters() if p.requires_grad
+        ]
        optimizer = torch.optim.AdamW(
            params=[
-                {'params': trainable_params},
+                {"params": trainable_params},
            ],
            lr=self.hparams.learning_rate,
            weight_decay=self.hparams.weight_decay,
@@ -324,13 +413,13 @@ class Pipeline(LightningModule):
                return float(current_step) / float(max(1, warmup_steps))
            else:
                # Linear decay from learning_rate to 0
-                progress = float(current_step - warmup_steps) / float(max(1, max_steps - warmup_steps))
+                progress = float(current_step - warmup_steps) / float(
+                    max(1, max_steps - warmup_steps)
+                )
                return max(0.0, 1.0 - progress)

        lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
-            optimizer,
-            lr_lambda,
-            last_epoch=-1
+            optimizer, lr_lambda, last_epoch=-1
        )
        return [optimizer], lr_scheduler

@@ -362,10 +451,17 @@ class Pipeline(LightningModule):
            # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
            # In practice, we sample the random variable u from a normal distribution u ∼ N (u; m, s)
            # and map it through the standard logistic function
-            u = torch.normal(mean=self.hparams.logit_mean, std=self.hparams.logit_std, size=(bsz, ), device="cpu")
+            u = torch.normal(
+                mean=self.hparams.logit_mean,
+                std=self.hparams.logit_std,
+                size=(bsz,),
+                device="cpu",
+            )
            u = torch.nn.functional.sigmoid(u)
            indices = (u * self.scheduler.config.num_train_timesteps).long()
-            indices = torch.clamp(indices, 0, self.scheduler.config.num_train_timesteps - 1)
+            indices = torch.clamp(
+                indices, 0, self.scheduler.config.num_train_timesteps - 1
+            )
            timesteps = self.scheduler.timesteps[indices].to(device)

        return timesteps
@@ -394,7 +490,9 @@ class Pipeline(LightningModule):
        timesteps = self.get_timestep(bsz, device)

        # Add noise according to flow matching.
-        sigmas = self.get_sd3_sigmas(timesteps=timesteps, device=device, n_dim=target_image.ndim, dtype=dtype)
+        sigmas = self.get_sd3_sigmas(
+            timesteps=timesteps, device=device, n_dim=target_image.ndim, dtype=dtype
+        )
        noisy_image = sigmas * noise + (1.0 - sigmas) * target_image

        # This is the flow-matching target for vanilla SD3.
@@ -431,7 +529,11 @@ class Pipeline(LightningModule):
        # Compute loss. Only calculate loss where chunk_mask is 1 and there is no padding
        # N x T x 64
        # N x T -> N x c x W x T
-        mask = attention_mask.unsqueeze(1).unsqueeze(1).expand(-1, target_image.shape[1], target_image.shape[2], -1)
+        mask = (
+            attention_mask.unsqueeze(1)
+            .unsqueeze(1)
+            .expand(-1, target_image.shape[1], target_image.shape[2], -1)
+        )

        selected_model_pred = (model_pred * mask).reshape(bsz, -1).contiguous()
        selected_target = (target * mask).reshape(bsz, -1).contiguous()
@@ -443,11 +545,19 @@ class Pipeline(LightningModule):

        prefix = "train"

-        self.log(f"{prefix}/denoising_loss", loss, on_step=True, on_epoch=False, prog_bar=True)
+        self.log(
+            f"{prefix}/denoising_loss",
+            loss,
+            on_step=True,
+            on_epoch=False,
+            prog_bar=True,
+        )

        total_proj_loss = 0.0
        for k, v in proj_losses:
-            self.log(f"{prefix}/{k}_loss", v, on_step=True, on_epoch=False, prog_bar=True)
+            self.log(
+                f"{prefix}/{k}_loss", v, on_step=True, on_epoch=False, prog_bar=True
+            )
            total_proj_loss += v

        if len(proj_losses) > 0:
@@ -459,7 +569,13 @@ class Pipeline(LightningModule):
        # Log learning rate if scheduler exists
        if self.lr_schedulers() is not None:
            learning_rate = self.lr_schedulers().get_last_lr()[0]
-            self.log(f"{prefix}/learning_rate", learning_rate, on_step=True, on_epoch=False, prog_bar=True)
+            self.log(
+                f"{prefix}/learning_rate",
+                learning_rate,
+                on_step=True,
+                on_epoch=False,
+                prog_bar=True,
+            )
        # with torch.autograd.detect_anomaly():
        #     self.manual_backward(loss)
        return loss
@@ -496,18 +612,35 @@ class Pipeline(LightningModule):
        )

        frame_length = int(duration * 44100 / 512 / 8)
-        timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps=infer_steps, device=device, timesteps=None)
+        timesteps, num_inference_steps = retrieve_timesteps(
+            scheduler, num_inference_steps=infer_steps, device=device, timesteps=None
+        )

-        target_latents = randn_tensor(shape=(bsz, 8, 16, frame_length), generator=random_generators, device=device, dtype=dtype)
+        target_latents = randn_tensor(
+            shape=(bsz, 8, 16, frame_length),
+            generator=random_generators,
+            device=device,
+            dtype=dtype,
+        )
        attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
        if do_classifier_free_guidance:
            attention_mask = torch.cat([attention_mask] * 2, dim=0)
-            encoder_text_hidden_states = torch.cat([encoder_text_hidden_states, torch.zeros_like(encoder_text_hidden_states)], 0)
+            encoder_text_hidden_states = torch.cat(
+                [
+                    encoder_text_hidden_states,
+                    torch.zeros_like(encoder_text_hidden_states),
+                ],
+                0,
+            )
            text_attention_mask = torch.cat([text_attention_mask] * 2, dim=0)

-            speaker_embds = torch.cat([speaker_embds, torch.zeros_like(speaker_embds)], 0)
+            speaker_embds = torch.cat(
+                [speaker_embds, torch.zeros_like(speaker_embds)], 0
+            )

-            lyric_token_ids = torch.cat([lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0)
+            lyric_token_ids = torch.cat(
+                [lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0
+            )
            lyric_mask = torch.cat([lyric_mask, torch.zeros_like(lyric_mask)], 0)

        momentum_buffer = MomentumBuffer()
@@ -515,7 +648,9 @@ class Pipeline(LightningModule):
        for i, t in tqdm(enumerate(timesteps), total=num_inference_steps):
            # expand the latents if we are doing classifier free guidance
            latents = target_latents
-            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = (
+                torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            )
            timestep = t.expand(latent_model_input.shape[0])
            noise_pred = self.transformers(
                hidden_states=latent_model_input,
@@ -537,7 +672,13 @@ class Pipeline(LightningModule):
                    momentum_buffer=momentum_buffer,
                )

-            target_latents = scheduler.step(model_output=noise_pred, timestep=t, sample=target_latents, return_dict=False, omega=omega_scale)[0]
+            target_latents = scheduler.step(
+                model_output=noise_pred,
+                timestep=t,
+                sample=target_latents,
+                return_dict=False,
+                omega=omega_scale,
+            )[0]

        return target_latents

@@ -582,7 +723,9 @@ class Pipeline(LightningModule):
        )

        audio_lengths = batch["wav_lengths"]
-        sr, pred_wavs = self.dcae.decode(pred_latents, audio_lengths=audio_lengths, sr=48000)
+        sr, pred_wavs = self.dcae.decode(
+            pred_latents, audio_lengths=audio_lengths, sr=48000
+        )
        return {
            "target_wavs": batch["target_wavs"],
            "pred_wavs": pred_wavs,
@@ -603,7 +746,12 @@ class Pipeline(LightningModule):

    def plot_step(self, batch, batch_idx):
        global_step = self.global_step
-        if global_step % self.hparams.every_plot_step != 0 or self.local_rank != 0 or torch.distributed.get_rank() != 0 or torch.cuda.current_device() != 0:
+        if (
+            global_step % self.hparams.every_plot_step != 0
+            or self.local_rank != 0
+            or torch.distributed.get_rank() != 0
+            or torch.cuda.current_device() != 0
+        ):
            return
        results = self.predict_step(batch)

@@ -615,7 +763,9 @@ class Pipeline(LightningModule):
        sr = results["sr"]
        seeds = results["seeds"]
        i = 0
-        for key, target_wav, pred_wav, prompt, candidate_lyric_chunk, seed in zip(keys, target_wavs, pred_wavs, prompts, candidate_lyric_chunks, seeds):
+        for key, target_wav, pred_wav, prompt, candidate_lyric_chunk, seed in zip(
+            keys, target_wavs, pred_wavs, prompts, candidate_lyric_chunks, seeds
+        ):
            key = key
            prompt = prompt
            lyric = self.construct_lyrics(candidate_lyric_chunk)
@@ -624,9 +774,15 @@ class Pipeline(LightningModule):
            save_dir = f"{log_dir}/eval_results/step_{self.global_step}"
            if not os.path.exists(save_dir):
                os.makedirs(save_dir, exist_ok=True)
-            torchaudio.save(f"{save_dir}/target_wav_{key}_{i}.flac", target_wav.float().cpu(), sr)
-            torchaudio.save(f"{save_dir}/pred_wav_{key}_{i}.flac", pred_wav.float().cpu(), sr)
-            with open(f"{save_dir}/key_prompt_lyric_{key}_{i}.txt", "w", encoding="utf-8") as f:
+            torchaudio.save(
+                f"{save_dir}/target_wav_{key}_{i}.wav", target_wav.float().cpu(), sr
+            )
+            torchaudio.save(
+                f"{save_dir}/pred_wav_{key}_{i}.wav", pred_wav.float().cpu(), sr
+            )
+            with open(
+                f"{save_dir}/key_prompt_lyric_{key}_{i}.txt", "w", encoding="utf-8"
+            ) as f:
                f.write(key_prompt_lyric)
            i += 1

@@ -642,11 +798,14 @@ def main(args):
        checkpoint_dir=args.checkpoint_dir,
    )
    checkpoint_callback = ModelCheckpoint(
-        monitor=None, every_n_train_steps=args.every_n_train_steps, save_top_k=-1,
+        monitor=None,
+        every_n_train_steps=args.every_n_train_steps,
+        save_top_k=-1,
    )
    # add datetime str to version
    logger_callback = TensorBoardLogger(
-        version=datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + args.exp_name, save_dir=args.logger_dir
+        version=datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + args.exp_name,
+        save_dir=args.logger_dir,
    )
    trainer = Trainer(
        accelerator="gpu",
				`@@ -0,0 +1 @@`
				`# This file intentionally left blank for Python to recognize the directory as a package.`