From 0feca8a820f09c9dde55da5ab06da4d8781d7824 Mon Sep 17 00:00:00 2001
From: rkfg <rkfg@rkfg.me>
Date: Tue, 13 May 2025 14:39:38 +0300
Subject: [PATCH] Add format selector

---
 acestep/pipeline_ace_step.py | 13 +++++--------
 acestep/ui/components.py     |  6 ++++++
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/acestep/pipeline_ace_step.py b/acestep/pipeline_ace_step.py
index 4799b20..a0aff31 100644
--- a/acestep/pipeline_ace_step.py
+++ b/acestep/pipeline_ace_step.py
@@ -1533,17 +1533,13 @@ class ACEStepPipeline:
             base_path = f"./outputs"
             ensure_directory_exists(base_path)
             output_path_wav = (
-                f"{base_path}/output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.wav"
+                f"{base_path}/output_{time.strftime('%Y%m%d%H%M%S')}_{idx}."+format
             )
         else:
             ensure_directory_exists(os.path.dirname(save_path))
             if os.path.isdir(save_path):
-                logger.info(
-                    f"Provided save_path '{save_path}' is a directory. Appending timestamped filename."
-                )
-                output_path_wav = os.path.join(
-                    save_path, f"output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.wav"
-                )
+                logger.info(f"Provided save_path '{save_path}' is a directory. Appending timestamped filename.")
+                output_path_wav = os.path.join(save_path, f"output_{time.strftime('%Y%m%d%H%M%S')}_{idx}."+format)
             else:
                 output_path_wav = save_path
 
@@ -1582,6 +1578,7 @@ class ACEStepPipeline:
 
     def __call__(
         self,
+        format: str = "wav",
         audio_duration: float = 60.0,
         prompt: str = None,
         lyrics: str = None,
@@ -1616,7 +1613,6 @@ class ACEStepPipeline:
         edit_n_max: float = 1.0,
         edit_n_avg: int = 1,
         save_path: str = None,
-        format: str = "wav",
         batch_size: int = 1,
         debug: bool = False,
     ):
@@ -1832,6 +1828,7 @@ class ACEStepPipeline:
         }
 
         input_params_json = {
+            "format": format,
             "lora_name_or_path": lora_name_or_path,
             "task": task,
             "prompt": prompt if task != "edit" else edit_target_prompt,
diff --git a/acestep/ui/components.py b/acestep/ui/components.py
index 2590c22..9378913 100644
--- a/acestep/ui/components.py
+++ b/acestep/ui/components.py
@@ -116,6 +116,7 @@ def create_text2music_ui(
                     info="-1 means random duration (30 ~ 240).",
                     scale=9,
                 )
+                format = gr.Dropdown(choices=["mp3", "ogg", "flac", "wav"], value="wav", label="Format")
                 sample_bnt = gr.Button("Sample", variant="secondary", scale=1)
 
             # audio2audio
@@ -319,6 +320,7 @@ def create_text2music_ui(
 
                 def retake_process_func(json_data, retake_variance, retake_seeds):
                     return text2music_process_func(
+                        json_data["format"],
                         json_data["audio_duration"],
                         json_data["prompt"],
                         json_data["lyrics"],
@@ -447,6 +449,7 @@ def create_text2music_ui(
                         src_audio_path = json_data["audio_path"]
 
                     return text2music_process_func(
+                        format.value,
                         json_data["audio_duration"],
                         prompt,
                         lyrics,
@@ -620,6 +623,7 @@ def create_text2music_ui(
                         edit_lyrics = lyrics
 
                     return text2music_process_func(
+                        format.value,
                         json_data["audio_duration"],
                         prompt,
                         lyrics,
@@ -766,6 +770,7 @@ def create_text2music_ui(
                     repaint_start = -left_extend_length
                     repaint_end = json_data["audio_duration"] + right_extend_length
                     return text2music_process_func(
+                        format.value,
                         json_data["audio_duration"],
                         prompt,
                         lyrics,
@@ -937,6 +942,7 @@ def create_text2music_ui(
     text2music_bnt.click(
         fn=text2music_process_func,
         inputs=[
+            format,
             audio_duration,
             prompt,
             lyrics,