add edit

pipeline_ace_step.py

@@ -56,7 +56,7 @@ REPO_ID = "ACE-Step/ACE-Step-v1-3.5B"
 # class ACEStepPipeline(DiffusionPipeline):
 class ACEStepPipeline:
 
-    def __init__(self, checkpoint_dir=None, device_id=0, dtype="bfloat16", text_encoder_checkpoint_path=None, **kwargs):
+    def __init__(self, checkpoint_dir=None, device_id=0, dtype="bfloat16", text_encoder_checkpoint_path=None, torch_compile=False, **kwargs):
         # check checkpoint dir exist
         if not checkpoint_dir:
             checkpoint_dir = os.path.join(os.path.dirname(__file__), "checkpoints")
@@ -72,6 +72,7 @@ class ACEStepPipeline:
         self.dtype = torch.float16 if dtype == "bfloat16" else torch.float32
         self.device = device
         self.loaded = False
+        self.torch_compile = torch_compile
 
     def load_checkpoint(self, checkpoint_dir=None):
         device = self.device
@@ -105,9 +106,10 @@ class ACEStepPipeline:
         self.loaded = True
 
         # compile
-        self.music_dcae = torch.compile(self.music_dcae)
+        if self.torch_compile:
-        self.ace_step_transformer = torch.compile(self.ace_step_transformer)
+            self.music_dcae = torch.compile(self.music_dcae)
-        self.text_encoder_model = torch.compile(self.text_encoder_model)
+            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):
         inputs = self.text_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=text_max_length)
@@ -215,6 +217,249 @@ class ACEStepPipeline:
                 print("tokenize error", e, "for line", line, "major_language", lang)
         return lyric_token_idx
 
+    def calc_v(
+        self,
+        zt_src,
+        zt_tar,
+        t,
+        encoder_text_hidden_states,
+        text_attention_mask,
+        target_encoder_text_hidden_states,
+        target_text_attention_mask,
+        speaker_embds,
+        target_speaker_embeds,
+        lyric_token_ids,
+        lyric_mask,
+        target_lyric_token_ids,
+        target_lyric_mask,
+        do_classifier_free_guidance=False,
+        guidance_scale=1.0,
+        target_guidance_scale=1.0,
+        cfg_type="apg",
+        attention_mask=None,
+        momentum_buffer=None,
+        momentum_buffer_tar=None,
+        return_src_pred=True
+    ):
+        noise_pred_src = None
+        if return_src_pred:
+            src_latent_model_input = torch.cat([zt_src, zt_src]) if do_classifier_free_guidance else zt_src
+            timestep = t.expand(src_latent_model_input.shape[0])
+            # source
+            noise_pred_src = self.ace_step_transformer(
+                hidden_states=src_latent_model_input,
+                attention_mask=attention_mask,
+                encoder_text_hidden_states=encoder_text_hidden_states,
+                text_attention_mask=text_attention_mask,
+                speaker_embeds=speaker_embds,
+                lyric_token_idx=lyric_token_ids,
+                lyric_mask=lyric_mask,
+                timestep=timestep,
+            ).sample
+
+            if do_classifier_free_guidance:
+                noise_pred_with_cond_src, noise_pred_uncond_src = noise_pred_src.chunk(2)
+                if cfg_type == "apg":
+                    noise_pred_src = apg_forward(
+                        pred_cond=noise_pred_with_cond_src,
+                        pred_uncond=noise_pred_uncond_src,
+                        guidance_scale=guidance_scale,
+                        momentum_buffer=momentum_buffer,
+                    )
+                elif cfg_type == "cfg":
+                    noise_pred_src = cfg_forward(
+                        cond_output=noise_pred_with_cond_src,
+                        uncond_output=noise_pred_uncond_src,
+                        cfg_strength=guidance_scale,
+                    )
+
+        tar_latent_model_input = torch.cat([zt_tar, zt_tar]) if do_classifier_free_guidance else zt_tar
+        timestep = t.expand(tar_latent_model_input.shape[0])
+        # target
+        noise_pred_tar = self.ace_step_transformer(
+            hidden_states=tar_latent_model_input,
+            attention_mask=attention_mask,
+            encoder_text_hidden_states=target_encoder_text_hidden_states,
+            text_attention_mask=target_text_attention_mask,
+            speaker_embeds=target_speaker_embeds,
+            lyric_token_idx=target_lyric_token_ids,
+            lyric_mask=target_lyric_mask,
+            timestep=timestep,
+        ).sample
+
+        if do_classifier_free_guidance:
+            noise_pred_with_cond_tar, noise_pred_uncond_tar = noise_pred_tar.chunk(2)
+            if cfg_type == "apg":
+                noise_pred_tar = apg_forward(
+                    pred_cond=noise_pred_with_cond_tar,
+                    pred_uncond=noise_pred_uncond_tar,
+                    guidance_scale=target_guidance_scale,
+                    momentum_buffer=momentum_buffer_tar,
+                )
+            elif cfg_type == "cfg":
+                noise_pred_tar = cfg_forward(
+                    cond_output=noise_pred_with_cond_tar,
+                    uncond_output=noise_pred_uncond_tar,
+                    cfg_strength=target_guidance_scale,
+                )
+        return noise_pred_src, noise_pred_tar
+
+    @torch.no_grad()
+    def flowedit_diffusion_process(
+        self,
+        encoder_text_hidden_states,
+        text_attention_mask,
+        speaker_embds,
+        lyric_token_ids,
+        lyric_mask,
+        target_encoder_text_hidden_states,
+        target_text_attention_mask,
+        target_speaker_embeds,
+        target_lyric_token_ids,
+        target_lyric_mask,
+        src_latents,
+        random_generators=None,
+        infer_steps=60,
+        guidance_scale=15.0,
+        n_min=0,
+        n_max=1.0,
+        n_avg=1,
+    ):
+
+        do_classifier_free_guidance = True
+        if guidance_scale == 0.0 or guidance_scale == 1.0:
+            do_classifier_free_guidance = False
+
+        target_guidance_scale = guidance_scale
+        device = encoder_text_hidden_states.device
+        dtype = encoder_text_hidden_states.dtype
+        bsz = encoder_text_hidden_states.shape[0]
+
+        scheduler = FlowMatchEulerDiscreteScheduler(
+            num_train_timesteps=1000,
+            shift=3.0,
+        )
+
+        T_steps = infer_steps
+        frame_length = src_latents.shape[-1]
+        attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
+        
+        timesteps, T_steps = retrieve_timesteps(scheduler, T_steps, device, timesteps=None)
+
+        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)
+            text_attention_mask = torch.cat([text_attention_mask] * 2, dim=0)
+
+            target_encoder_text_hidden_states = torch.cat([target_encoder_text_hidden_states, torch.zeros_like(target_encoder_text_hidden_states)], 0)
+            target_text_attention_mask = torch.cat([target_text_attention_mask] * 2, dim=0)
+
+            speaker_embds = torch.cat([speaker_embds, torch.zeros_like(speaker_embds)], 0)
+            target_speaker_embeds = torch.cat([target_speaker_embeds, torch.zeros_like(target_speaker_embeds)], 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)
+
+            target_lyric_token_ids = torch.cat([target_lyric_token_ids, torch.zeros_like(target_lyric_token_ids)], 0)
+            target_lyric_mask = torch.cat([target_lyric_mask, torch.zeros_like(target_lyric_mask)], 0)
+
+        momentum_buffer = MomentumBuffer()
+        momentum_buffer_tar = MomentumBuffer()
+        x_src = src_latents
+        zt_edit = x_src.clone()
+        xt_tar = None
+        n_min = int(infer_steps * n_min)
+        n_max = int(infer_steps * n_max)
+
+        logger.info("flowedit start from {} to {}".format(n_min, n_max))
+
+        for i, t in tqdm(enumerate(timesteps), total=T_steps):
+
+            if i < n_min:
+                continue
+
+            t_i = t/1000
+
+            if i+1 < len(timesteps): 
+                t_im1 = (timesteps[i+1])/1000
+            else:
+                t_im1 = torch.zeros_like(t_i).to(t_i.device)
+
+            if i < n_max:
+                # Calculate the average of the V predictions
+                V_delta_avg = torch.zeros_like(x_src)
+                for k in range(n_avg):
+                    fwd_noise = randn_tensor(shape=x_src.shape, generator=random_generators, device=device, dtype=dtype)
+
+                    zt_src = (1 - t_i) * x_src + (t_i) * fwd_noise
+
+                    zt_tar = zt_edit + zt_src - x_src
+
+                    Vt_src, Vt_tar = self.calc_v(
+                        zt_src=zt_src,
+                        zt_tar=zt_tar,
+                        t=t,
+                        encoder_text_hidden_states=encoder_text_hidden_states,
+                        text_attention_mask=text_attention_mask,
+                        target_encoder_text_hidden_states=target_encoder_text_hidden_states,
+                        target_text_attention_mask=target_text_attention_mask,
+                        speaker_embds=speaker_embds,
+                        target_speaker_embeds=target_speaker_embeds,
+                        lyric_token_ids=lyric_token_ids,
+                        lyric_mask=lyric_mask,
+                        target_lyric_token_ids=target_lyric_token_ids,
+                        target_lyric_mask=target_lyric_mask,
+                        do_classifier_free_guidance=do_classifier_free_guidance,
+                        guidance_scale=guidance_scale,
+                        target_guidance_scale=target_guidance_scale,
+                        attention_mask=attention_mask,
+                        momentum_buffer=momentum_buffer
+                    )
+                    V_delta_avg += (1 / n_avg) * (Vt_tar - Vt_src) # - (hfg-1)*( x_src))
+
+                # propagate direct ODE
+                zt_edit = zt_edit.to(torch.float32)
+                zt_edit = zt_edit + (t_im1 - t_i) * V_delta_avg
+                zt_edit = zt_edit.to(V_delta_avg.dtype)
+            else: # i >= T_steps-n_min # regular sampling for last n_min steps
+                if i == n_max:
+                    fwd_noise = randn_tensor(shape=x_src.shape, generator=random_generators, device=device, dtype=dtype)
+                    scheduler._init_step_index(t)
+                    sigma = scheduler.sigmas[scheduler.step_index]
+                    xt_src = sigma * fwd_noise + (1.0 - sigma) * x_src
+                    xt_tar = zt_edit + xt_src - x_src
+                
+                _, Vt_tar = self.calc_v(
+                    zt_src=None,
+                    zt_tar=xt_tar,
+                    t=t,
+                    encoder_text_hidden_states=encoder_text_hidden_states,
+                    text_attention_mask=text_attention_mask,
+                    target_encoder_text_hidden_states=target_encoder_text_hidden_states,
+                    target_text_attention_mask=target_text_attention_mask,
+                    speaker_embds=speaker_embds,
+                    target_speaker_embeds=target_speaker_embeds,
+                    lyric_token_ids=lyric_token_ids,
+                    lyric_mask=lyric_mask,
+                    target_lyric_token_ids=target_lyric_token_ids,
+                    target_lyric_mask=target_lyric_mask,
+                    do_classifier_free_guidance=do_classifier_free_guidance,
+                    guidance_scale=guidance_scale,
+                    target_guidance_scale=target_guidance_scale,
+                    attention_mask=attention_mask,
+                    momentum_buffer_tar=momentum_buffer_tar,
+                    return_src_pred=False,
+                )
+                
+                dtype = Vt_tar.dtype
+                xt_tar = xt_tar.to(torch.float32)
+                prev_sample = xt_tar + (t_im1 - t_i) * Vt_tar
+                prev_sample = prev_sample.to(dtype) 
+                xt_tar = prev_sample
+        
+        target_latents = zt_edit if xt_tar is None else xt_tar
+        return target_latents
 
     @torch.no_grad()
     def text2music_diffusion_process(
@@ -593,6 +838,11 @@ class ACEStepPipeline:
         repaint_start: int = 0,
         repaint_end: int = 0,
         src_audio_path: str = None,
+        edit_target_prompt: str = None,
+        edit_target_lyrics: str = None,
+        edit_n_min: float = 0.0,
+        edit_n_max: float = 1.0,
+        edit_n_avg: int = 1,
         save_path: str = None,
         format: str = "flac",
         batch_size: int = 1,
@@ -653,40 +903,76 @@ class ACEStepPipeline:
             repaint_end = audio_duration
         
         src_latents = None
-        if task == "repaint":
+        if src_audio_path is not None:
-            assert src_audio_path is not None, "src_audio_path is required for repaint task"
+            assert src_audio_path is not None and task in ("repaint", "edit"), "src_audio_path is required for repaint task"
             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)
 
-        target_latents = self.text2music_diffusion_process(
+        if task == "edit":
-            duration=audio_duration,
+            texts = [edit_target_prompt]
-            encoder_text_hidden_states=encoder_text_hidden_states,
+            target_encoder_text_hidden_states, target_text_attention_mask = self.get_text_embeddings(texts, self.device)
-            text_attention_mask=text_attention_mask,
+            target_encoder_text_hidden_states = target_encoder_text_hidden_states.repeat(batch_size, 1, 1)
-            speaker_embds=speaker_embeds,
+            target_text_attention_mask = target_text_attention_mask.repeat(batch_size, 1)
-            lyric_token_ids=lyric_token_idx,
+
-            lyric_mask=lyric_mask,
+            target_lyric_token_idx = torch.tensor([0]).repeat(batch_size, 1).to(self.device).long()
-            guidance_scale=guidance_scale,
+            target_lyric_mask = torch.tensor([0]).repeat(batch_size, 1).to(self.device).long()
-            omega_scale=omega_scale,
+            if len(edit_target_lyrics) > 0:
-            infer_steps=infer_step,
+                target_lyric_token_idx = self.tokenize_lyrics(edit_target_lyrics, debug=True)
-            random_generators=random_generators,
+                target_lyric_mask = [1] * len(target_lyric_token_idx)
-            scheduler_type=scheduler_type,
+                target_lyric_token_idx = torch.tensor(target_lyric_token_idx).unsqueeze(0).to(self.device).repeat(batch_size, 1)
-            cfg_type=cfg_type,
+                target_lyric_mask = torch.tensor(target_lyric_mask).unsqueeze(0).to(self.device).repeat(batch_size, 1)
-            guidance_interval=guidance_interval,
+
-            guidance_interval_decay=guidance_interval_decay,
+            target_speaker_embeds = speaker_embeds.clone()
-            min_guidance_scale=min_guidance_scale,
+
-            oss_steps=oss_steps,
+            target_latents = self.flowedit_diffusion_process(
-            encoder_text_hidden_states_null=encoder_text_hidden_states_null,
+                encoder_text_hidden_states=encoder_text_hidden_states,
-            use_erg_lyric=use_erg_lyric,
+                text_attention_mask=text_attention_mask,
-            use_erg_diffusion=use_erg_diffusion,
+                speaker_embds=speaker_embeds,
-            retake_random_generators=retake_random_generators,
+                lyric_token_ids=lyric_token_idx,
-            retake_variance=retake_variance,
+                lyric_mask=lyric_mask,
-            add_retake_noise=add_retake_noise,
+                target_encoder_text_hidden_states=target_encoder_text_hidden_states,
-            guidance_scale_text=guidance_scale_text,
+                target_text_attention_mask=target_text_attention_mask,
-            guidance_scale_lyric=guidance_scale_lyric,
+                target_speaker_embeds=target_speaker_embeds,
-            repaint_start=repaint_start,
+                target_lyric_token_ids=target_lyric_token_idx,
-            repaint_end=repaint_end,
+                target_lyric_mask=target_lyric_mask,
-            src_latents=src_latents,
+                src_latents=src_latents,
-        )
+                random_generators=random_generators,
+                infer_steps=infer_step,
+                guidance_scale=guidance_scale,
+                n_min=edit_n_min,
+                n_max=edit_n_max,
+                n_avg=edit_n_avg,
+            )
+        else:
+            target_latents = self.text2music_diffusion_process(
+                duration=audio_duration,
+                encoder_text_hidden_states=encoder_text_hidden_states,
+                text_attention_mask=text_attention_mask,
+                speaker_embds=speaker_embeds,
+                lyric_token_ids=lyric_token_idx,
+                lyric_mask=lyric_mask,
+                guidance_scale=guidance_scale,
+                omega_scale=omega_scale,
+                infer_steps=infer_step,
+                random_generators=random_generators,
+                scheduler_type=scheduler_type,
+                cfg_type=cfg_type,
+                guidance_interval=guidance_interval,
+                guidance_interval_decay=guidance_interval_decay,
+                min_guidance_scale=min_guidance_scale,
+                oss_steps=oss_steps,
+                encoder_text_hidden_states_null=encoder_text_hidden_states_null,
+                use_erg_lyric=use_erg_lyric,
+                use_erg_diffusion=use_erg_diffusion,
+                retake_random_generators=retake_random_generators,
+                retake_variance=retake_variance,
+                add_retake_noise=add_retake_noise,
+                guidance_scale_text=guidance_scale_text,
+                guidance_scale_lyric=guidance_scale_lyric,
+                repaint_start=repaint_start,
+                repaint_end=repaint_end,
+                src_latents=src_latents,
+            )
 
         end_time = time.time()
         diffusion_time_cost = end_time - start_time
@@ -730,6 +1016,14 @@ class ACEStepPipeline:
             "retake_variance": retake_variance,
             "guidance_scale_text": guidance_scale_text,
             "guidance_scale_lyric": guidance_scale_lyric,
+            "repaint_start": repaint_start,
+            "repaint_end": repaint_end,
+            "edit_n_min": edit_n_min,
+            "edit_n_max": edit_n_max,
+            "edit_n_avg": edit_n_avg,        
+            "src_audio_path": src_audio_path,
+            "edit_target_prompt": edit_target_prompt,
+            "edit_target_lyrics": edit_target_lyrics,
         }
         # save input_params_json
         for output_audio_path in output_paths:

ui/components.py

@@ -98,8 +98,8 @@ def create_text2music_ui(
         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_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, info="Seed for the retake.")
+                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")
                 
@@ -138,8 +138,8 @@ 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", info="Variance for the retake. 0.0 means no variance. 1.0 means full variance.")
+                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, info="Seed for the retake.")
+                retake_seeds = gr.Textbox(label="retake 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")
@@ -250,7 +250,139 @@ def create_text2music_ui(
                     outputs=repaint_outputs + [repaint_input_params_json],
                 )
             with gr.Tab("edit"):
-                pass
+                edit_prompt = gr.Textbox(lines=2, label="Edit Tags", max_lines=4)
+                edit_lyrics = gr.Textbox(lines=9, label="Edit Lyrics", max_lines=13)
+
+                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.8, 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.8
+                        n_max = 1.0
+                    elif edit_type == "remix":
+                        n_min = 0.2
+                        n_max = 0.4
+                    return n_min, n_max
+                
+                edit_type.change(
+                    edit_type_change_func,
+                    inputs=[edit_type],
+                    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.change(
+                    fn=lambda x: gr.update(visible=x == "upload", elem_id="edit_source_audio_upload"),
+                    inputs=[edit_source],
+                    outputs=[edit_source_audio_upload],
+                )
+
+                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,
+                    edit_source,
+                    edit_source_audio_upload,
+                    prompt,
+                    lyrics,
+                    edit_prompt,
+                    edit_lyrics,
+                    edit_n_min,
+                    edit_n_max,
+                    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,
+                ):
+                    if edit_source == "upload":
+                        src_audio_path = edit_source_audio_upload
+                        json_data = text2music_json_data
+                    elif edit_source == "text2music":
+                        json_data = text2music_json_data
+                        src_audio_path = json_data["audio_path"]
+                    elif edit_source == "last_edit":
+                        json_data = edit_input_params_json
+                        src_audio_path = json_data["audio_path"]
+
+                    if not edit_prompt:
+                        edit_prompt = prompt
+                    if not edit_lyrics:
+                        edit_lyrics = lyrics
+
+                    return text2music_process_func(
+                        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,
+                        task="edit",
+                        src_audio_path=src_audio_path,
+                        edit_target_prompt=edit_prompt,
+                        edit_target_lyrics=edit_lyrics,
+                        edit_n_min=edit_n_min,
+                        edit_n_max=edit_n_max
+                    )
+                
+                edit_bnt.click(
+                    fn=edit_process_func,
+                    inputs=[
+                        input_params_json,
+                        edit_input_params_json,
+                        edit_source,
+                        edit_source_audio_upload,
+                        prompt,
+                        lyrics,
+                        edit_prompt,
+                        edit_lyrics,
+                        edit_n_min,
+                        edit_n_max,
+                        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,
+                    ],
+                    outputs=edit_outputs + [edit_input_params_json],
+                )
 
         def sample_data():
             json_data = sample_data_func()