rename to ace step

README.md

@@ -1,15 +1,14 @@
-# Fusic
+# ACE-Step
 
-Installation instructions for Fusic.
+# Installation
-## Requirements
 ```bash
-conda create -n fusic python==3.10
+conda create -n ace_step python==3.10
-conda activate fusic
+conda activate ace_step
 pip install -r requirements.txt
 conda install ffmpeg
 ```
 
-## Usage
+# Usage
 ```bash
 python app.py
 ```

app.py

@@ -1,6 +1,6 @@
 import argparse
 parser = argparse.ArgumentParser()
-parser.add_argument("--checkpoint_path", type=str, default="./checkpoints")
+parser.add_argument("--checkpoint_path", type=str, default="")
 parser.add_argument("--port", type=int, default=7865)
 parser.add_argument("--device_id", type=int, default=0)
 parser.add_argument("--share", action='store_true', default=False)
@@ -14,13 +14,13 @@ os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
 
 
 from ui.components import create_main_demo_ui
-from pipeline_fusic import FusicPipeline
+from pipeline_ace_step import ACEStepPipeline
 from data_sampler import DataSampler
 
 
 def main(args):
  
-    model_demo = FusicPipeline(
+    model_demo = ACEStepPipeline(
         checkpoint_dir=args.checkpoint_path,
         dtype="bfloat16" if args.bf16 else "float32"
     )

hf_download.py

@@ -246,4 +246,4 @@ def download_repo_multi_part(repo: str, save_path: str, headers):
 
 
 if __name__ == "__main__":
-    download_repo("timedomain/fusic_v1", "checkpoints_new")
+    download_repo("timedomain/ace_step_v1", "checkpoints")

models/ace_step_transformer.py

@@ -153,7 +153,7 @@ class Transformer2DModelOutput(BaseOutput):
     proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None
 
 
-class FusicTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
     _supports_gradient_checkpointing = True
 
     @register_to_config

pipeline_ace_step.py

@@ -21,7 +21,7 @@ from hf_download import download_repo
 
 from language_segmentation import LangSegment
 from music_dcae.music_dcae_pipeline import MusicDCAE
-from models.fusic_transformer import FusicTransformer2DModel
+from models.ace_step_transformer import ACEStepTransformer2DModel
 from models.lyrics_utils.lyric_tokenizer import VoiceBpeTokenizer
 from apg_guidance import apg_forward, MomentumBuffer, cfg_forward, cfg_zero_star, cfg_double_condition_forward
 import torchaudio
@@ -44,22 +44,22 @@ def ensure_directory_exists(directory):
         os.makedirs(directory)
 
 
-REPO_ID = "timedomain/fusic_v1"
+REPO_ID = "ACE-Step/ACE-Step-v1-3.5B"
 
 
-# class FusicPipeline(DiffusionPipeline):
+# class ACEStepPipeline(DiffusionPipeline):
-class FusicPipeline:
+class ACEStepPipeline:
 
     def __init__(self, checkpoint_dir=None, device_id=0, dtype="bfloat16", text_encoder_checkpoint_path=None, **kwargs):
         # check checkpoint dir exist
-        if checkpoint_dir is None:
+        if not checkpoint_dir:
             checkpoint_dir = os.path.join(os.path.dirname(__file__), "checkpoints")
-            if not os.path.exists(checkpoint_dir):
+        if not os.path.exists(checkpoint_dir):
-                # huggingface download
+            # huggingface download
-                download_repo(
+            download_repo(
-                    repo_id=REPO_ID,
+                repo_id=REPO_ID,
-                    save_path=checkpoint_dir
+                save_path=checkpoint_dir
-                )
+            )
 
         self.checkpoint_dir = checkpoint_dir
         device = torch.device(f"cuda:{device_id}") if torch.cuda.is_available() else torch.device("cpu")
@@ -74,9 +74,9 @@ class FusicPipeline:
         self.music_dcae = MusicDCAE(dcae_checkpoint_path=dcae_checkpoint_path, vocoder_checkpoint_path=vocoder_checkpoint_path)
         self.music_dcae.to(device).eval().to(self.dtype)
 
-        fusic_checkpoint_path = os.path.join(checkpoint_dir, "fusic_transformer")
+        ace_step_checkpoint_path = os.path.join(checkpoint_dir, "ace_step_transformer")
-        self.fusic_transformer = FusicTransformer2DModel.from_pretrained(fusic_checkpoint_path)
+        self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(ace_step_checkpoint_path)
-        self.fusic_transformer.to(device).eval().to(self.dtype)
+        self.ace_step_transformer.to(device).eval().to(self.dtype)
 
         lang_segment = LangSegment()
 
@@ -100,7 +100,7 @@ class FusicPipeline:
 
         # compile
         self.music_dcae = torch.compile(self.music_dcae)
-        self.fusic_transformer = torch.compile(self.fusic_transformer)
+        self.ace_step_transformer = torch.compile(self.ace_step_transformer)
         self.text_encoder_model = torch.compile(self.text_encoder_model)
 
     def get_text_embeddings(self, texts, device, text_max_length=256):
@@ -305,10 +305,10 @@ class FusicPipeline:
                 return output
             
             for i in range(l_min, l_max):
-                handler = self.fusic_transformer.lyric_encoder.encoders[i].self_attn.linear_q.register_forward_hook(hook)
+                handler = self.ace_step_transformer.lyric_encoder.encoders[i].self_attn.linear_q.register_forward_hook(hook)
                 handlers.append(handler)
         
-            encoder_hidden_states, encoder_hidden_mask = self.fusic_transformer.encode(**inputs)
+            encoder_hidden_states, encoder_hidden_mask = self.ace_step_transformer.encode(**inputs)
             
             for hook in handlers:
                 hook.remove()
@@ -316,7 +316,7 @@ class FusicPipeline:
             return encoder_hidden_states
 
         # P(speaker, text, lyric)
-        encoder_hidden_states, encoder_hidden_mask = self.fusic_transformer.encode(
+        encoder_hidden_states, encoder_hidden_mask = self.ace_step_transformer.encode(
             encoder_text_hidden_states,
             text_attention_mask,
             speaker_embds,
@@ -338,7 +338,7 @@ class FusicPipeline:
             )
         else:
             # P(null_speaker, null_text, null_lyric)
-            encoder_hidden_states_null, _ = self.fusic_transformer.encode(
+            encoder_hidden_states_null, _ = self.face_step_transformer.encode(
                 torch.zeros_like(encoder_text_hidden_states),
                 text_attention_mask,
                 torch.zeros_like(speaker_embds),
@@ -362,7 +362,7 @@ class FusicPipeline:
                 )
             # P(null_speaker, text, no_lyric)
             else:
-                encoder_hidden_states_no_lyric, _ = self.fusic_transformer.encode(
+                encoder_hidden_states_no_lyric, _ = self.ace_step_transformer.encode(
                     encoder_text_hidden_states,
                     text_attention_mask,
                     torch.zeros_like(speaker_embds),
@@ -378,12 +378,12 @@ class FusicPipeline:
                 return output
             
             for i in range(l_min, l_max):
-                handler = self.fusic_transformer.transformer_blocks[i].attn.to_q.register_forward_hook(hook)
+                handler = self.ace_step_transformer.transformer_blocks[i].attn.to_q.register_forward_hook(hook)
                 handlers.append(handler)
-                handler = self.fusic_transformer.transformer_blocks[i].cross_attn.to_q.register_forward_hook(hook)
+                handler = self.ace_step_transformer.transformer_blocks[i].cross_attn.to_q.register_forward_hook(hook)
                 handlers.append(handler)
 
-            sample = self.fusic_transformer.decode(hidden_states=hidden_states, timestep=timestep, **inputs).sample
+            sample = self.ace_step_transformer.decode(hidden_states=hidden_states, timestep=timestep, **inputs).sample
             
             for hook in handlers:
                 hook.remove()
@@ -409,7 +409,7 @@ class FusicPipeline:
                 timestep = t.expand(latent_model_input.shape[0])
                 output_length = latent_model_input.shape[-1]
                 # P(x|speaker, text, lyric)
-                noise_pred_with_cond = self.fusic_transformer.decode(
+                noise_pred_with_cond = self.ace_step_transformer.decode(
                     hidden_states=latent_model_input,
                     attention_mask=attention_mask,
                     encoder_hidden_states=encoder_hidden_states,
@@ -420,7 +420,7 @@ class FusicPipeline:
 
                 noise_pred_with_only_text_cond = None
                 if do_double_condition_guidance and encoder_hidden_states_no_lyric is not None:
-                    noise_pred_with_only_text_cond = self.fusic_transformer.decode(
+                    noise_pred_with_only_text_cond = self.ace_step_transformer.decode(
                         hidden_states=latent_model_input,
                         attention_mask=attention_mask,
                         encoder_hidden_states=encoder_hidden_states_no_lyric,
@@ -442,7 +442,7 @@ class FusicPipeline:
                         },
                     )
                 else:
-                    noise_pred_uncond = self.fusic_transformer.decode(
+                    noise_pred_uncond = self.ace_step_transformer.decode(
                         hidden_states=latent_model_input,
                         attention_mask=attention_mask,
                         encoder_hidden_states=encoder_hidden_states_null,
@@ -485,7 +485,7 @@ class FusicPipeline:
             else:
                 latent_model_input = latents
                 timestep = t.expand(latent_model_input.shape[0])
-                noise_pred = self.fusic_transformer.decode(
+                noise_pred = self.ace_step_transformer.decode(
                     hidden_states=latent_model_input,
                     attention_mask=attention_mask,
                     encoder_hidden_states=encoder_hidden_states,
@@ -514,7 +514,7 @@ class FusicPipeline:
     def save_wav_file(self, target_wav, idx, save_path=None, sample_rate=48000, format="flac"):
         if save_path is None:
             logger.warning("save_path is None, using default path ./outputs/")
-            base_path = f"./outputs/"
+            base_path = f"./outputs"
             ensure_directory_exists(base_path)
         else:
             base_path = save_path

requirements.txt

@@ -19,4 +19,4 @@ num2words==0.5.14
 spacy==3.8.4
 accelerate==1.6.0
 cutlet
-'fugashi[unidic-lite]'
+fugashi[unidic-lite]

ui/components.py

@@ -63,15 +63,15 @@ def create_text2music_ui(
 ):
     with gr.Row():
         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=180, 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, placeholder=TAG_PLACEHOLDER, info="Support tags, descriptions, and scene. Use commas to separate different tags.")
+            prompt = gr.Textbox(lines=2, label="Tags", max_lines=4, placeholder=TAG_PLACEHOLDER, 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, placeholder=LYRIC_PLACEHOLDER, 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=True):
+            with gr.Accordion("Basic Settings", open=False):
                 infer_step = gr.Slider(minimum=1, maximum=1000, step=1, value=60, 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=5.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")
@@ -93,14 +93,14 @@ def create_text2music_ui(
                 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(variant="primary")
+            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", info="Variance for the retake. 0.0 means no variance. 1.0 means full variance.")
                 retake_seeds = gr.Textbox(label="retake seeds (default None)", placeholder="", value=None, info="Seed for the retake.")
-                retake_bnt = gr.Button(variant="primary")
+                retake_bnt = gr.Button("Retake", variant="primary")
                 retake_outputs, retake_input_params_json = create_output_ui("Retake")
                 
                 def retake_process_func(json_data, retake_variance, retake_seeds):
@@ -219,13 +219,12 @@ def create_main_demo_ui(
     sample_data_func=dump_func,
 ):
     with gr.Blocks(
-        title="FusicModel 1.0 DEMO",
+        title="ACE-Step Model 1.0 DEMO",
     ) as demo:
         gr.Markdown(
             """
-            <h1 style="text-align: center;">FusicModel 1.0 DEMO</h1>
+            <h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
-            """
+        """)
-        )
 
         with gr.Tab("text2music"):
             create_text2music_ui(