update

app.py

@@ -38,6 +38,7 @@ def infer_music(lrc, ref_audio_path, steps, file_type, max_frames=2048):
     style_prompt = get_style_prompt(muq, ref_audio_path)
     negative_style_prompt = get_negative_style_prompt(device)
     latent_prompt = get_reference_latent(device, max_frames)
+    print(">0")
     generated_song = inference(cfm_model=cfm, 
                                vae_model=vae, 
                                cond=latent_prompt, 
@@ -52,6 +53,7 @@ def infer_music(lrc, ref_audio_path, steps, file_type, max_frames=2048):
                                )
     torch.cuda.empty_cache()
     gc.collect()
+    print(">4")
 
     return generated_song
 

diffrhythm/infer/infer.py

@@ -78,6 +78,7 @@ def decode_audio(latents, vae_model, chunked=False, overlap=32, chunk_size=128):
 def inference(cfm_model, vae_model, cond, text, duration, style_prompt, negative_style_prompt, steps, sway_sampling_coef, start_time, file_type):
 
     with torch.inference_mode():
+        print(">1")
         generated, _ = cfm_model.sample(
             cond=cond,
             text=text,
@@ -93,14 +94,19 @@ def inference(cfm_model, vae_model, cond, text, duration, style_prompt, negative
         gc.collect()
 
         
+        print(">2")
         generated = generated.to(torch.float32)
+        print(">3")
         latent = generated.transpose(1, 2) # [b d t]
-        output = decode_audio(latent, vae_model, chunked=False)
+        print(">4")
+        output = decode_audio(latent, vae_model, chunked=True)
+        print(">5")
 
         del latent, generated
         torch.cuda.empty_cache()
         gc.collect()
 
+        print(">6")
 
         # Rearrange audio batch to a single sequence
         output = rearrange(output, "b d n -> d (b n)")