Update app.py

app.py

@@ -98,17 +98,6 @@ DEMUCS_MODELS = [
     'hdemucs_mmi.yaml',
 ]
 
-class CustomProgressBar:
-    def __init__(self, total=1.0):
-        self.total = total
-        self.current = 0.0
-        self.desc = ""
-
-    def update(self, value, desc):
-        self.current = value
-        self.desc = desc
-        print(f"Progress: {self.current * 100:.2f}% - {self.desc}")
-
 def print_message(input_file, model_name):
     """Prints information about the audio separation process."""
     base_name = os.path.splitext(os.path.basename(input_file))[0]
@@ -146,7 +135,7 @@ def prepare_output_dir(input_file, output_dir):
         raise
     return out_dir
 
-def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=CustomProgressBar()):
+def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using Roformer model."""
     print_message(audio, model_key)
     model = ROFORMER_MODELS[model_key]
@@ -168,13 +157,13 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
             }
         )
 
-        progress.update(0.2, "Model loaded")
+        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
 
-        progress.update(0.7, "Audio separated")
+        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio)
 
-        progress.update(0.9, "Stems renamed")
+        progress(0.9, desc="Stems renamed...")
         stems = rename_stems(audio, out_dir, separation, out_format)
 
         return stems[0], stems[1]
@@ -182,7 +171,7 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
         logging.error(f"Roformer separation failed: {e}")
         return None, None
 
-def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=CustomProgressBar()):
+def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_shift, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using MDX23C model."""
     print_message(audio, model)
     try:
@@ -203,13 +192,13 @@ def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_s
             }
         )
 
-        progress.update(0.2, "Model loaded")
+        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
 
-        progress.update(0.7, "Audio separated")
+        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio)
 
-        progress.update(0.9, "Stems renamed")
+        progress(0.9, desc="Stems renamed...")
         stems = rename_stems(audio, out_dir, separation, out_format)
 
         return stems[0], stems[1]
@@ -217,7 +206,7 @@ def mdx23c_separator(audio, model, seg_size, override_seg_size, overlap, pitch_s
         logging.error(f"MDX23C separation failed: {e}")
         return None, None
 
-def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=CustomProgressBar()):
+def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using MDX-NET model."""
     print_message(audio, model)
     try:
@@ -238,13 +227,13 @@ def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_di
             }
         )
 
-        progress.update(0.2, "Model loaded")
+        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
 
-        progress.update(0.7, "Audio separated")
+        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio)
 
-        progress.update(0.9, "Stems renamed")
+        progress(0.9, desc="Stems renamed...")
         stems = rename_stems(audio, out_dir, separation, out_format)
 
         return stems[0], stems[1]
@@ -252,7 +241,7 @@ def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_di
         logging.error(f"MDX-NET separation failed: {e}")
         return None, None
 
-def vr_separator(audio, model, window_size, aggression, tta, post_process, post_process_threshold, high_end_process, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=CustomProgressBar()):
+def vr_separator(audio, model, window_size, aggression, tta, post_process, post_process_threshold, high_end_process, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using VR ARCH model."""
     print_message(audio, model)
     try:
@@ -275,13 +264,13 @@ def vr_separator(audio, model, window_size, aggression, tta, post_process, post_
             }
         )
 
-        progress.update(0.2, "Model loaded")
+        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
 
-        progress.update(0.7, "Audio separated")
+        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio)
 
-        progress.update(0.9, "Stems renamed")
+        progress(0.9, desc="Stems renamed...")
         stems = rename_stems(audio, out_dir, separation, out_format)
 
         return stems[0], stems[1]
@@ -289,7 +278,7 @@ def vr_separator(audio, model, window_size, aggression, tta, post_process, post_
         logging.error(f"VR ARCH separation failed: {e}")
         return None, None
 
-def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=CustomProgressBar()):
+def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress(track_tqdm=True)):
     """Separate audio using Demucs model."""
     print_message(audio, model)
     try:
@@ -309,13 +298,13 @@ def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled,
             }
         )
 
-        progress.update(0.2, "Model loaded")
+        progress(0.2, desc="Model loaded...")
         separator.load_model(model_filename=model)
 
-        progress.update(0.7, "Audio separated")
+        progress(0.7, desc="Audio separated...")
         separation = separator.separate(audio)
 
-        progress.update(0.9, "Stems renamed")
+        progress(0.9, desc="Stems renamed...")
         stems = rename_stems(audio, out_dir, separation, out_format)
 
         return stems[0], stems[1], stems[2], stems[3]