support audio2audio

pipeline_ace_step.py

@@ -2,7 +2,7 @@ import random
 import time
 import os
 import re
-import spaces
+# import spaces
 import torch
 import torch.nn as nn
 from loguru import logger
@@ -522,6 +522,27 @@ class ACEStepPipeline:
         target_latents = zt_edit if xt_tar is None else xt_tar
         return target_latents
 
+    def add_latents_noise(
+        self,
+        gt_latents,
+        variance,
+        noise,
+        scheduler,
+    ):
+
+        bsz = gt_latents.shape[0]
+        u = torch.tensor([variance] * bsz, dtype=gt_latents.dtype)
+        indices = (u * scheduler.config.num_train_timesteps).long()
+        timesteps = scheduler.timesteps.unsqueeze(1).to(gt_latents.dtype)
+        indices = indices.to(timesteps.device).to(gt_latents.dtype).unsqueeze(1)
+        nearest_idx = torch.argmin(torch.cdist(indices, timesteps), dim=1)
+        sigma = scheduler.sigmas[nearest_idx].flatten().to(gt_latents.device).to(gt_latents.dtype)
+        while len(sigma.shape) < gt_latents.ndim:
+            sigma = sigma.unsqueeze(-1)
+        noisy_image = sigma * noise + (1.0 - sigma) * gt_latents
+        init_timestep = indices[0]
+        return noisy_image, init_timestep
+
     @torch.no_grad()
     def text2music_diffusion_process(
         self,
@@ -554,6 +575,9 @@ class ACEStepPipeline:
         repaint_start=0,
         repaint_end=0,
         src_latents=None,
+        audio2audio_enable=False,
+        ref_audio_strength=0.5,
+        ref_latents=None,
     ):
 
         logger.info("cfg_type: {}, guidance_scale: {}, omega_scale: {}".format(cfg_type, guidance_scale, omega_scale))
@@ -585,6 +609,9 @@ class ACEStepPipeline:
         if src_latents is not None:
             frame_length = src_latents.shape[-1]
 
+        if ref_latents is not None:
+            frame_length = ref_latents.shape[-1]
+
         if len(oss_steps) > 0:
             infer_steps = max(oss_steps)
             scheduler.set_timesteps
@@ -680,6 +707,10 @@ class ACEStepPipeline:
                 zt_edit = x0.clone()
                 z0 = target_latents
 
+        init_timestep = 1000
+        if audio2audio_enable and ref_latents is not None:
+            target_latents, init_timestep = self.add_latents_noise(gt_latents=ref_latents, variance=(1-ref_audio_strength), noise=target_latents, scheduler=scheduler)
+
         attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
         
         # guidance interval
@@ -783,7 +814,10 @@ class ACEStepPipeline:
             return sample
     
         for i, t in tqdm(enumerate(timesteps), total=num_inference_steps):
-            
+
+            if t > init_timestep:
+                continue
+
             if is_repaint:
                 if i < n_min:
                     continue
@@ -955,7 +989,7 @@ class ACEStepPipeline:
         latents, _ = self.music_dcae.encode(input_audio, sr=sr)
         return latents
 
-    @spaces.GPU
+    # @spaces.GPU
     def __call__(
         self,
         audio_duration: float = 60.0,
@@ -976,6 +1010,9 @@ class ACEStepPipeline:
         oss_steps: str = None,
         guidance_scale_text: float = 0.0,
         guidance_scale_lyric: float = 0.0,
+        audio2audio_enable: bool = False,
+        ref_audio_strength: float = 0.5,
+        ref_audio_input: str = None,
         retake_seeds: list = None,
         retake_variance: float = 0.5,
         task: str = "text2music",
@@ -995,6 +1032,9 @@ class ACEStepPipeline:
 
         start_time = time.time()
 
+        if audio2audio_enable and ref_audio_input is not None:
+            task = "audio2audio"
+
         if not self.loaded:
             logger.warning("Checkpoint not loaded, loading checkpoint...")
             self.load_checkpoint(self.checkpoint_dir)
@@ -1053,6 +1093,14 @@ class ACEStepPipeline:
             assert os.path.exists(src_audio_path), f"src_audio_path {src_audio_path} does not exist"
             src_latents = self.infer_latents(src_audio_path)
 
+        ref_latents = None
+        if ref_audio_input is not None and audio2audio_enable:
+            assert ref_audio_input is not None, "ref_audio_input is required for audio2audio task"
+            assert os.path.exists(
+                ref_audio_input
+            ), f"ref_audio_input {ref_audio_input} does not exist"
+            ref_latents = self.infer_latents(ref_audio_input)
+
         if task == "edit":
             texts = [edit_target_prompt]
             target_encoder_text_hidden_states, target_text_attention_mask = self.get_text_embeddings(texts, self.device)
@@ -1117,6 +1165,9 @@ class ACEStepPipeline:
                 repaint_start=repaint_start,
                 repaint_end=repaint_end,
                 src_latents=src_latents,
+                audio2audio_enable=audio2audio_enable,
+                ref_audio_strength=ref_audio_strength,
+                ref_latents=ref_latents,
             )
 
         end_time = time.time()
@@ -1169,6 +1220,9 @@ class ACEStepPipeline:
             "src_audio_path": src_audio_path,
             "edit_target_prompt": edit_target_prompt,
             "edit_target_lyrics": edit_target_lyrics,
+            "audio2audio_enable": audio2audio_enable,
+            "ref_audio_strength": ref_audio_strength,
+            "ref_audio_input": ref_audio_input,
         }
         # save input_params_json
         for output_audio_path in output_paths:

ui/components.py

@@ -71,6 +71,32 @@ def create_text2music_ui(
                 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)
 
+            # audio2audio
+            audio2audio_enable = gr.Checkbox(label="Enable Audio2Audio", value=False, info="Check to enable Audio-to-Audio generation using a reference audio.", elem_id="audio2audio_checkbox")
+            ref_audio_input = gr.Audio(type="filepath", label="Reference Audio (for Audio2Audio)", visible=False, elem_id="ref_audio_input", show_download_button=True)
+            ref_audio_strength = gr.Slider(
+                label="Refer audio strength",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.01,
+                value=0.5,
+                elem_id="ref_audio_strength",
+                visible=False,
+                interactive=True,
+            )
+
+            def toggle_ref_audio_visibility(is_checked):
+                return (
+                    gr.update(visible=is_checked, elem_id="ref_audio_input"),
+                    gr.update(visible=is_checked, elem_id="ref_audio_strength"),
+                )
+
+            audio2audio_enable.change(
+                fn=toggle_ref_audio_visibility,
+                inputs=[audio2audio_enable],
+                outputs=[ref_audio_input, ref_audio_strength],
+            )
+
             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")
 
@@ -533,6 +559,21 @@ def create_text2music_ui(
                 ", ".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["audio2audio_enable"]
+                    if "audio2audio_enable" in json_data
+                    else False
+                ),
+                (
+                    json_data["ref_audio_strength"]
+                    if "ref_audio_strength" in json_data
+                    else 0.5
+                ),
+                (
+                    json_data["ref_audio_input"]
+                    if "ref_audio_input" in json_data
+                    else None
+                ),
             )
 
         sample_bnt.click(
@@ -556,6 +597,9 @@ def create_text2music_ui(
                 oss_steps,
                 guidance_scale_text,
                 guidance_scale_lyric,
+                audio2audio_enable,
+                ref_audio_strength,
+                ref_audio_input,
             ],
         )
 
@@ -580,6 +624,9 @@ def create_text2music_ui(
             oss_steps,
             guidance_scale_text,
             guidance_scale_lyric,
+            audio2audio_enable,
+            ref_audio_strength,
+            ref_audio_input,
         ], outputs=outputs + [input_params_json]
     )