[fix] fix some bugs

.gitignore

@@ -1,3 +1,16 @@
+*.pt
+*.ckpt
+*.onnx
+t5_g2p_model/
+embeddings/
+checkpoints/
+
+val_images/
+val_audios/
+lightning_logs/
+lightning_logs_/
+train_images/
+train_audios/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -7,6 +20,7 @@ __pycache__/
 *.so
 
 # Distribution / packaging
+.idea/
 .Python
 build/
 develop-eggs/
@@ -99,17 +113,15 @@ ipython_config.py
 #   This is especially recommended for binary packages to ensure reproducibility, and is more
 #   commonly ignored for libraries.
 #   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-poetry.lock
+#poetry.lock
 
 # pdm
 #   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
 #pdm.lock
 #   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
 #   in version control.
-#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+#   https://pdm.fming.dev/#use-with-ide
 .pdm.toml
-.pdm-python
-.pdm-build/
 
 # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
 __pypackages__/
@@ -160,5 +172,31 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
-
-checkpoints/
+*.txt
+!requirements.txt
+*.log
+*.flac
+minio_config.yaml
+.history/*
+__pycache__/*
+train.log
+*.mp3
+*.tar.gz
+__pycache__/
+demo_examples/
+nohup.out
+test_results/*
+nohup.out
+text_audio_align/*
+remote/*
+MG2P/*
+audio_getter.py
+refiner_loss_debug/
+outputs/*
+!outputs/
+save_checkpoint.ipynb
+repos/*
+app_demo.py
+ui/components_demo.py
+data_sampler_demo.py
+pipeline_ace_step_demo.py

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [2025] Timedomain Inc. and stepfun
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

apg_guidance.py

@@ -17,7 +17,10 @@ def project(
     dims=[-1, -2],
 ):
     dtype = v0.dtype
-    v0, v1 = v0.double(), v1.double()
+    if v0.device.type == "mps":
+        v0, v1 = v0.float(), v1.float()
+    else:
+        v0, v1 = v0.double(), v1.double()
     v1 = torch.nn.functional.normalize(v1, dim=dims)
     v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
     v0_orthogonal = v0 - v0_parallel
@@ -53,6 +56,7 @@ def apg_forward(
 def cfg_forward(cond_output, uncond_output, cfg_strength):
     return uncond_output + cfg_strength * (cond_output - uncond_output)
 
+
 def cfg_double_condition_forward(
     cond_output,
     uncond_output,

music_dcae/music_dcae_pipeline.py

@@ -3,7 +3,6 @@ import torch
 from diffusers import AutoencoderDC
 import torchaudio
 import torchvision.transforms as transforms
-import torchaudio
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.loaders import FromOriginalModelMixin
 from diffusers.configuration_utils import ConfigMixin, register_to_config
@@ -30,7 +29,7 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
 
         if source_sample_rate is None:
             source_sample_rate = 48000
-        
+
         self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
 
         self.transform = transforms.Compose([
@@ -95,29 +94,21 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
     def decode(self, latents, audio_lengths=None, sr=None):
         latents = latents / self.scale_factor + self.shift_factor
 
-        mels = []
+        pred_wavs = []
 
         for latent in latents:
-            mel = self.dcae.decoder(latent.unsqueeze(0))
-            mels.append(mel)
-        mels = torch.cat(mels, dim=0)
-
-        mels = mels * 0.5 + 0.5
-        mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
-        bsz, channels, num_mel, mel_width = mels.shape
-        pred_wavs = []
-        for i in range(bsz):
-            mel = mels[i]
-            wav = self.vocoder.decode(mel).squeeze(1)
+            mels = self.dcae.decoder(latent.unsqueeze(0))
+            mels = mels * 0.5 + 0.5
+            mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
+            wav = self.vocoder.decode(mels[0]).squeeze(1)
+
+            if sr is not None:
+                resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
+                wav = resampler(wav)
+            else:
+                sr = 44100
             pred_wavs.append(wav)
 
-        pred_wavs = torch.stack(pred_wavs)
-
-        if sr is not None:
-            resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
-            pred_wavs = [resampler(wav) for wav in pred_wavs]
-        else:
-            sr = 44100
         if audio_lengths is not None:
             pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
         return sr, pred_wavs

pipeline_ace_step.py

@@ -11,6 +11,7 @@ import json
 import math
 from huggingface_hub import hf_hub_download
 
+# from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 from schedulers.scheduling_flow_match_heun_discrete import FlowMatchHeunDiscreteScheduler
 from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
@@ -29,6 +30,7 @@ torch.backends.cudnn.benchmark = False
 torch.set_float32_matmul_precision('high')
 torch.backends.cudnn.deterministic = True
 torch.backends.cuda.matmul.allow_tf32 = True
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
 SUPPORT_LANGUAGES = {
@@ -54,7 +56,7 @@ REPO_ID = "ACE-Step/ACE-Step-v1-3.5B"
 class ACEStepPipeline:
 
     def __init__(self, checkpoint_dir=None, device_id=0, dtype="bfloat16", text_encoder_checkpoint_path=None, persistent_storage_path=None, torch_compile=False, **kwargs):
-        if checkpoint_dir is None:
+        if not checkpoint_dir:
             if persistent_storage_path is None:
                 checkpoint_dir = os.path.join(os.path.dirname(__file__), "checkpoints")
             else:
@@ -63,7 +65,11 @@ class ACEStepPipeline:
 
         self.checkpoint_dir = checkpoint_dir
         device = torch.device(f"cuda:{device_id}") if torch.cuda.is_available() else torch.device("cpu")
+        if device.type == "cpu" and torch.backends.mps.is_available():
+            device = torch.device("mps")
         self.dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float32
+        if device.type == "mps" and self.dtype == torch.bfloat16:
+            self.dtype = torch.float16
         self.device = device
         self.loaded = False
         self.torch_compile = torch_compile
@@ -620,6 +626,7 @@ class ACEStepPipeline:
                 repaint_mask[:, :, :, repaint_start_frame:repaint_end_frame] = 1.0
                 repaint_noise = torch.cos(retake_variance) * target_latents + torch.sin(retake_variance) * retake_latents
                 repaint_noise = torch.where(repaint_mask == 1.0, repaint_noise, target_latents)
+                zt_edit = x0.clone()
                 z0 = repaint_noise
             elif is_extend:
                 to_right_pad_gt_latents = None
@@ -669,9 +676,8 @@ class ACEStepPipeline:
                     padd_list.append(retake_latents[:, :, :, -right_pad_frame_length:])
                 target_latents = torch.cat(padd_list, dim=-1)
                 assert target_latents.shape[-1] == x0.shape[-1], f"{target_latents.shape=} {x0.shape=}"
-
-            zt_edit = x0.clone()
-            z0 = target_latents
+                zt_edit = x0.clone()
+                z0 = target_latents
 
         attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
         
@@ -774,7 +780,7 @@ class ACEStepPipeline:
                 hook.remove()
             
             return sample
-
+    
         for i, t in tqdm(enumerate(timesteps), total=num_inference_steps):
             
             if is_repaint:

ui/components.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import librosa
 
 
 TAG_PLACEHOLDER = "funk, pop, soul, rock, melodic, guitar, drums, bass, keyboard, percussion, 105 BPM, energetic, upbeat, groovy, vibrant, dynamic"
@@ -67,25 +68,25 @@ def create_text2music_ui(
                 # 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)
                 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.\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=False):
-                infer_step = gr.Slider(minimum=1, maximum=1000, step=1, value=60, label="Infer Steps", 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=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")
-                guidance_scale_lyric = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=1.5, label="Guidance Scale Lyric", interactive=True)
+                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")
-            
+
             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.")
-        
+
                 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)")
@@ -102,7 +103,7 @@ def create_text2music_ui(
                 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")
-                
+
                 def retake_process_func(json_data, retake_variance, retake_seeds):
                     return text2music_process_func(
                         json_data["audio_duration"],
@@ -143,7 +144,7 @@ def create_text2music_ui(
                 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.change(
                     fn=lambda x: gr.update(visible=x == "upload", elem_id="repaint_source_audio_upload"),
@@ -153,7 +154,7 @@ def create_text2music_ui(
 
                 repaint_bnt = gr.Button("Repaint", variant="primary")
                 repaint_outputs, repaint_input_params_json = create_output_ui("Repaint")
-                
+
                 def repaint_process_func(
                     text2music_json_data,
                     repaint_json_data,
@@ -183,7 +184,10 @@ def create_text2music_ui(
                 ):
                     if repaint_source == "upload":
                         src_audio_path = repaint_source_audio_upload
-                        json_data = text2music_json_data
+                        audio_duration = librosa.get_duration(filename=src_audio_path)
+                        json_data = {
+                            "audio_duration": audio_duration
+                        }
                     elif repaint_source == "text2music":
                         json_data = text2music_json_data
                         src_audio_path = json_data["audio_path"]
@@ -217,7 +221,7 @@ def create_text2music_ui(
                         repaint_end=repaint_end,
                         src_audio_path=src_audio_path,
                     )
-                
+
                 repaint_bnt.click(
                     fn=repaint_process_func,
                     inputs=[
@@ -253,11 +257,11 @@ def create_text2music_ui(
                 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)
-                
+
                 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":
                         n_min = 0.6
@@ -266,7 +270,7 @@ def create_text2music_ui(
                         n_min = 0.2
                         n_max = 0.4
                     return n_min, n_max
-                
+
                 edit_type.change(
                     edit_type_change_func,
                     inputs=[edit_type],
@@ -283,7 +287,7 @@ def create_text2music_ui(
 
                 edit_bnt = gr.Button("Edit", variant="primary")
                 edit_outputs, edit_input_params_json = create_output_ui("Edit")
-                
+
                 def edit_process_func(
                     text2music_json_data,
                     edit_input_params_json,
@@ -314,7 +318,10 @@ def create_text2music_ui(
                 ):
                     if edit_source == "upload":
                         src_audio_path = edit_source_audio_upload
-                        json_data = text2music_json_data
+                        audio_duration = librosa.get_duration(filename=src_audio_path)
+                        json_data = {
+                            "audio_duration": audio_duration
+                        }
                     elif edit_source == "text2music":
                         json_data = text2music_json_data
                         src_audio_path = json_data["audio_path"]
@@ -354,7 +361,7 @@ def create_text2music_ui(
                         edit_n_max=edit_n_max,
                         retake_seeds=retake_seeds,
                     )
-                
+
                 edit_bnt.click(
                     fn=edit_process_func,
                     inputs=[
@@ -392,7 +399,7 @@ def create_text2music_ui(
                 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.change(
                     fn=lambda x: gr.update(visible=x == "upload", elem_id="extend_source_audio_upload"),
@@ -402,7 +409,7 @@ def create_text2music_ui(
 
                 extend_bnt = gr.Button("Extend", variant="primary")
                 extend_outputs, extend_input_params_json = create_output_ui("Extend")
-                
+
                 def extend_process_func(
                     text2music_json_data,
                     extend_input_params_json,
@@ -431,11 +438,15 @@ def create_text2music_ui(
                 ):
                     if extend_source == "upload":
                         src_audio_path = extend_source_audio_upload
-                        json_data = text2music_json_data
+                        # get audio duration
+                        audio_duration = librosa.get_duration(filename=src_audio_path)
+                        json_data = {
+                            "audio_duration": audio_duration
+                        }
                     elif extend_source == "text2music":
                         json_data = text2music_json_data
                         src_audio_path = json_data["audio_path"]
-                    elif extend_source == "last_repaint":
+                    elif extend_source == "last_extend":
                         json_data = extend_input_params_json
                         src_audio_path = json_data["audio_path"]
 
@@ -467,7 +478,7 @@ def create_text2music_ui(
                         repaint_end=repaint_end,
                         src_audio_path=src_audio_path,
                     )
-                
+
                 extend_bnt.click(
                     fn=extend_process_func,
                     inputs=[
@@ -521,7 +532,7 @@ def create_text2music_ui(
                 json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
                 json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
             )
-    
+
         sample_bnt.click(
             sample_data,
             outputs=[