fix repaint and edit

pipeline_ace_step.py

@@ -69,7 +69,7 @@ class ACEStepPipeline:
 
         self.checkpoint_dir = checkpoint_dir
         device = torch.device(f"cuda:{device_id}") if torch.cuda.is_available() else torch.device("cpu")
-        self.dtype = torch.float16 if dtype == "bfloat16" else torch.float32
+        self.dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float32
         self.device = device
         self.loaded = False
         self.torch_compile = torch_compile
@@ -558,6 +558,8 @@ class ACEStepPipeline:
                 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)
                 z0 = repaint_noise
+                n_min = int(infer_steps * (1 - retake_variance))
+                zt_edit = src_latents.clone()
 
         attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
         
@@ -609,7 +611,7 @@ class ACEStepPipeline:
             )
         else:
             # P(null_speaker, null_text, null_lyric)
-            encoder_hidden_states_null, _ = self.face_step_transformer.encode(
+            encoder_hidden_states_null, _ = self.ace_step_transformer.encode(
                 torch.zeros_like(encoder_text_hidden_states),
                 text_attention_mask,
                 torch.zeros_like(speaker_embds),
@@ -660,9 +662,18 @@ class ACEStepPipeline:
                 hook.remove()
             
             return sample
-
     
         for i, t in tqdm(enumerate(timesteps), total=num_inference_steps):
+            
+            if is_repaint:
+                if i < n_min:
+                    continue
+                elif i == n_min:
+                    t_i = t / 1000
+                    zt_src = (1 - t_i) * src_latents + (t_i) * z0
+                    target_latents = zt_edit + zt_src - src_latents
+                    logger.info(f"repaint start from {n_min} add {t_i} level of noise")
+
             # expand the latents if we are doing classifier free guidance
             latents = target_latents
 
@@ -765,14 +776,22 @@ class ACEStepPipeline:
                     timestep=timestep,
                 ).sample
 
-            target_latents = scheduler.step(model_output=noise_pred, timestep=t, sample=target_latents, return_dict=False, omega=omega_scale)[0]
+            if is_repaint and i >= n_min:
-            if is_repaint:
+                t_i = t/1000
-                t_i = t / 1000
+                if i+1 < len(timesteps): 
-                x0 = src_latents
+                    t_im1 = (timesteps[i+1])/1000
-                xt = (1 - t_i) * x0 + t_i * z0
+                else:
-                target_latents = torch.where(repaint_mask == 1.0, target_latents, xt)
+                    t_im1 = torch.zeros_like(t_i).to(t_i.device)
+                dtype = noise_pred.dtype
+                target_latents = target_latents.to(torch.float32)
+                prev_sample = target_latents + (t_im1 - t_i) * noise_pred
+                prev_sample = prev_sample.to(dtype)
+                target_latents = prev_sample
+                zt_src = (1 - t_im1) * src_latents + (t_im1) * z0
+                target_latents = torch.where(repaint_mask == 1.0, target_latents, zt_src)
+            else:
+                target_latents = scheduler.step(model_output=noise_pred, timestep=t, sample=target_latents, return_dict=False, omega=omega_scale)[0]
 
-        
         return target_latents
 
     def latents2audio(self, latents, target_wav_duration_second=30, sample_rate=48000, save_path=None, format="flac"):
@@ -936,7 +955,7 @@ class ACEStepPipeline:
                 target_lyric_token_ids=target_lyric_token_idx,
                 target_lyric_mask=target_lyric_mask,
                 src_latents=src_latents,
-                random_generators=random_generators,
+                random_generators=retake_random_generators, # more diversity
                 infer_steps=infer_step,
                 guidance_scale=guidance_scale,
                 n_min=edit_n_min,
@@ -995,8 +1014,8 @@ class ACEStepPipeline:
 
         input_params_json = {
             "task": task,
-            "prompt": prompt,
+            "prompt": prompt if task != "edit" else edit_target_prompt,
-            "lyrics": lyrics,
+            "lyrics": lyrics if task != "edit" else edit_target_lyrics,
             "audio_duration": audio_duration,
             "infer_step": infer_step,
             "guidance_scale": guidance_scale,

ui/components.py

@@ -139,7 +139,7 @@ def create_text2music_ui(
                 )
             with gr.Tab("repainting"):
                 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)
+                retake_seeds = gr.Textbox(label="repaint 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")
@@ -252,14 +252,15 @@ def create_text2music_ui(
             with gr.Tab("edit"):
                 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.8, label="edit_n_min", interactive=True)
+                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.8
+                        n_min = 0.6
                         n_max = 1.0
                     elif edit_type == "remix":
                         n_min = 0.2
@@ -309,6 +310,7 @@ def create_text2music_ui(
                     oss_steps,
                     guidance_scale_text,
                     guidance_scale_lyric,
+                    retake_seeds,
                 ):
                     if edit_source == "upload":
                         src_audio_path = edit_source_audio_upload
@@ -349,7 +351,8 @@ def create_text2music_ui(
                         edit_target_prompt=edit_prompt,
                         edit_target_lyrics=edit_lyrics,
                         edit_n_min=edit_n_min,
-                        edit_n_max=edit_n_max
+                        edit_n_max=edit_n_max,
+                        retake_seeds=retake_seeds,
                     )
                 
                 edit_bnt.click(
@@ -380,6 +383,7 @@ def create_text2music_ui(
                         oss_steps,
                         guidance_scale_text,
                         guidance_scale_lyric,
+                        retake_seeds,
                     ],
                     outputs=edit_outputs + [edit_input_params_json],
                 )