update

app.py

@@ -9,6 +9,7 @@ from einops import rearrange
 import argparse
 import json
 import os
+import gc
 #import spaces
 from tqdm import tqdm
 import random
@@ -49,6 +50,9 @@ def infer_music(lrc, ref_audio_path, steps, file_type, max_frames=2048):
                                start_time=start_time,
                                file_type=file_type
                                )
+    torch.cuda.empty_cache()
+    gc.collect()
+
     return generated_song
 
 def R1_infer1(theme, tags_gen, language):

diffrhythm/infer/infer.py

@@ -9,6 +9,7 @@ import random
 import numpy as np
 import time
 import io
+import gc
 import pydub
 
 from diffrhythm.infer.infer_utils import (
@@ -88,11 +89,19 @@ def inference(cfm_model, vae_model, cond, text, duration, style_prompt, negative
             sway_sampling_coef=sway_sampling_coef,
             start_time=start_time
         )
+        torch.cuda.empty_cache()
+        gc.collect()
+
         
         generated = generated.to(torch.float32)
         latent = generated.transpose(1, 2) # [b d t]
         output = decode_audio(latent, vae_model, chunked=False)
 
+        del latent, generated
+        torch.cuda.empty_cache()
+        gc.collect()
+
+
         # Rearrange audio batch to a single sequence
         output = rearrange(output, "b d n -> d (b n)")
         output_tensor = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).cpu()
@@ -157,4 +166,4 @@ if __name__ == "__main__":
     
     output_path = os.path.join(output_dir, "output.wav")
     torchaudio.save(output_path, generated_song, sample_rate=44100)
-    
+