Merge pull request #4 from SUC-DriverOld/main

Add process callback and manage process bar

Author: baicai-1145

pymss/modules/vocal_remover/common_separator.py

@@ -35,7 +35,7 @@ def __init__(self, config):
         self.model_path = config.get("model_path")
         self.model_data = config.get("model_data")
         self.sample_rate = config.get("sample_rate")
-        self.callback = config.get("callback", None)
+        self.progress_callback = config.get("progress_callback", None)
 
         self.primary_stem_name = self.model_data.get("primary_stem", "primary_stem")
         self.secondary_stem_name = self.model_data.get("secondary_stem", "secondary_stem")

pymss/modules/vocal_remover/vr_separator.py

@@ -271,6 +271,8 @@ def execute(x_mag_pad, roi_size):
             write_pos = 0
             batch_starts = range(0, patches, self.batch_size)
             process_batches = tqdm(batch_starts, leave=False, desc="Processing VR batches") if self.debug else batch_starts
+            if self.progress_callback:
+                self.progress_callback(0, patches, "Processing VR batches")
             if self._use_mlx_full_forward(device):
                 import mlx.core as mx
 
@@ -280,8 +282,8 @@ def execute(x_mag_pad, roi_size):
                     pred = pred.astype(mx.float32).transpose(1, 2, 0, 3).reshape(pred.shape[1], pred.shape[2], -1)
                     mask_batches.append(pred)
                     write_pos += pred.shape[2]
-                    if self.callback:
-                        self.callback["progress"] = min(0.99 * (i / patches), 0.99)
+                    if self.progress_callback:
+                        self.progress_callback(min(i + self.batch_size, patches), patches, "Processing VR batches")
                 return mx.concatenate(mask_batches, axis=2)[:, :, :write_pos]
 
             with torch.inference_mode():
@@ -303,8 +305,8 @@ def execute(x_mag_pad, roi_size):
                         )
                     mask[:, :, write_pos:write_pos + pred.size(2)] = pred
                     write_pos += pred.size(2)
-                    if self.callback:
-                        self.callback["progress"] = min(0.99 * (i / patches), 0.99)
+                    if self.progress_callback:
+                        self.progress_callback(min(i + self.batch_size, patches), patches, "Processing VR batches")
             return mask[:, :, :write_pos]
 
         def adjust_aggr_torch(mask, is_non_accom_stem):

pymss/separator.py

@@ -362,6 +362,7 @@ def __init__(
             audio_params = {"wav_bit_depth": "FLOAT", "flac_bit_depth": "PCM_24", "mp3_bit_rate": "320k", "m4a_bit_rate": "192k", "m4a_aac_at_quality": 2},
             logger = None,
             debug = False,
+            progress_callback = None,
             inference_params = {
                 "batch_size": None,
                 "overlap_size": None,
@@ -390,6 +391,7 @@ def __init__(
         self.audio_params = audio_params
         self.logger = logger
         self.debug = debug
+        self.progress_callback = progress_callback
         self.inference_params = inference_params
 
         if self.debug:
@@ -497,7 +499,7 @@ def load_model(self):
                 "model_path": self.model_path,
                 "model_data": model_data,
                 "sample_rate": 44100,
-                "callback": None,
+                "progress_callback": self.progress_callback,
             }
             model = VRSeparator(common_config, config.inference)
             model.load_model()
@@ -745,7 +747,16 @@ def _separate(self, mix, pbar, stems=None):
         mix_orig = mix.copy()
         mix, norm_stats = _normalize_mix(mix, self.config.inference.get('normalize', False), self.logger)
         full_result = [
-            demix(self.config, self.model, track, self.device, pbar=pbar, model_type=self.model_type, source_indices=source_indices)
+            demix(
+                self.config,
+                self.model,
+                track,
+                self.device,
+                pbar=pbar,
+                model_type=self.model_type,
+                source_indices=source_indices,
+                progress_callback=self.progress_callback,
+            )
             for track in _tta_variants(mix, self.use_tta, self.logger)
         ]
 

pymss/utils.py

@@ -10,6 +10,45 @@
 from .config import load_config
 
 
+class _ProgressContext:
+    def __init__(self, pbar=False, total=1, callback=None, done=0, message="Processing audio chunks"):
+        self.enabled = bool(pbar or callback)
+        self.bar = None
+        self.callback = callback
+        self.done = done
+        self.total = total
+        self.message = message
+        if not self.enabled:
+            return
+        self.bar = tqdm(total=total, desc=message, leave=False) if pbar else None
+        self.total = int(self.total or 1)
+        self.done = int(self.done or 0)
+        self.emit()
+
+    def emit(self, done=None):
+        if not self.enabled:
+            return
+        if done is not None:
+            self.done = int(done)
+        if self.callback is None:
+            return
+        self.callback(min(self.done, self.total), self.total, self.message)
+
+    def update(self, amount):
+        if not self.enabled:
+            return
+        amount = int(amount)
+        if self.bar:
+            self.bar.update(amount)
+        self.done += amount
+        self.emit()
+
+    def close(self):
+        if not self.enabled:
+            return
+        if self.bar:
+            self.bar.close()
+
 def get_model_from_config(model_type, config_path, model_kwargs_override=None):
     model_kwargs_override = model_kwargs_override or {}
     config = load_config(config_path)
@@ -247,7 +286,18 @@ def _add_weighted_chunk(result, counter, chunk, window, start, length):
     counter[..., start:start + length] += window
 
 
-def _run_complete_chunks(model, mix, windows, result, counter, chunk_size, step, batch_size, progress_bar, source_indices=None):
+def _run_complete_chunks(
+    model,
+    mix,
+    windows,
+    result,
+    counter,
+    chunk_size,
+    step,
+    batch_size,
+    progress,
+    source_indices=None,
+):
     n_chunks = _complete_chunk_count(mix.shape[1], chunk_size, step)
     if n_chunks == 0:
         return 0
@@ -272,13 +322,25 @@ def _run_complete_chunks(model, mix, windows, result, counter, chunk_size, step,
             step,
             start_offset=batch_start * step,
         )
-        if progress_bar:
-            progress_bar.update(step * (batch_end - batch_start))
+        progress.update(step * (batch_end - batch_start))
 
     return n_complete
 
 
-def _run_tail_chunks(model, mix, starts, windows, result, counter, chunk_size, step, batch_size, first_chunk, progress_bar, source_indices=None):
+def _run_tail_chunks(
+    model,
+    mix,
+    starts,
+    windows,
+    result,
+    counter,
+    chunk_size,
+    step,
+    batch_size,
+    first_chunk,
+    progress,
+    source_indices=None,
+):
     for batch_start in range(first_chunk, len(starts), batch_size):
         batch_indices = range(batch_start, min(batch_start + batch_size, len(starts)))
         batch = [(_extract_chunk(mix, starts[idx], chunk_size), idx) for idx in batch_indices]
@@ -288,8 +350,7 @@ def _run_tail_chunks(model, mix, starts, windows, result, counter, chunk_size, s
             start = starts[idx]
             _add_weighted_chunk(result, counter, chunks[j], windows[idx], start, length)
 
-        if progress_bar:
-            progress_bar.update(step * len(batch_data))
+        progress.update(step * len(batch_data))
 
 
 def _finalize_overlap(result, counter, length_init, border):
@@ -449,7 +510,7 @@ def _can_demix_mlx_full(model, device):
     )
 
 
-def demix_track_mlx_full(config, model, mix, device, pbar=False, source_indices=None):
+def demix_track_mlx_full(config, model, mix, device, pbar=False, source_indices=None, progress_callback=None):
     import mlx.core as mx
 
     C = config.audio.chunk_size
@@ -463,7 +524,7 @@ def demix_track_mlx_full(config, model, mix, device, pbar=False, source_indices=
     starts, windows = _mlx_build_chunk_plan(mix.shape[1], C, step, fade_size)
     result = mx.zeros((_source_count(config, source_indices), mix.shape[0], mix.shape[1]), dtype=mx.float32)
     counter = mx.zeros((1, 1, mix.shape[1]), dtype=mx.float32)
-    progress_bar = tqdm(total=mix.shape[1], desc="Processing audio chunks", leave=False) if pbar else None
+    progress = _ProgressContext(pbar, mix.shape[1], progress_callback)
 
     for batch_start in range(0, len(starts), batch_size):
         batch_indices = range(batch_start, min(batch_start + batch_size, len(starts)))
@@ -473,18 +534,17 @@ def demix_track_mlx_full(config, model, mix, device, pbar=False, source_indices=
         for j, ((_, length), idx) in enumerate(batch):
             result, counter = _mlx_add_weighted_chunk(result, counter, chunks[j], windows[idx], starts[idx], length)
         mx.eval(result, counter)
-        if progress_bar:
-            progress_bar.update(step * len(batch))
+        progress.update(step * len(batch))
 
-    if progress_bar:
-        progress_bar.close()
+    progress.close()
+    progress.emit(mix.shape[1])
     return _sources_to_dict(config, _mlx_finalize_overlap(result, counter, length_init, border), source_indices)
 
 
 demix_track_mlx_roformer = demix_track_mlx_full
 
 
-def demix_track(config, model, mix, device, pbar=False, source_indices=None):
+def demix_track(config, model, mix, device, pbar=False, source_indices=None, progress_callback=None):
     C = config.audio.chunk_size
     source_indices = _normalize_source_indices(config, source_indices)
     step = _get_inference_step(config, C)
@@ -501,32 +561,51 @@ def demix_track(config, model, mix, device, pbar=False, source_indices=None):
     with _autocast(device, config.training.get('use_amp', True)):
         with torch.inference_mode():
             result, counter = _init_overlap_buffers(config, mix, device, use_complete_fast_path, source_indices)
-            progress_bar = tqdm(total=mix.shape[1], desc="Processing audio chunks", leave=False) if pbar else None
+            progress = _ProgressContext(pbar, mix.shape[1], progress_callback)
 
             with _model_source_context(model, source_indices):
                 complete_chunks = 0
                 if use_complete_fast_path:
                     complete_chunks = _run_complete_chunks(
-                        model, mix_device, chunk_windows, result, counter, C, step, batch_size, progress_bar, source_indices
+                        model,
+                        mix_device,
+                        chunk_windows,
+                        result,
+                        counter,
+                        C,
+                        step,
+                        batch_size,
+                        progress,
+                        source_indices,
                     )
 
                 _run_tail_chunks(
-                    model, mix_device, chunk_starts, chunk_windows, result, counter, C, step, batch_size,
-                    complete_chunks, progress_bar, source_indices
+                    model,
+                    mix_device,
+                    chunk_starts,
+                    chunk_windows,
+                    result,
+                    counter,
+                    C,
+                    step,
+                    batch_size,
+                    complete_chunks,
+                    progress,
+                    source_indices,
                 )
+                progress.emit(mix.shape[1])
 
 
-            if progress_bar:
-                progress_bar.close()
+            progress.close()
 
             estimated_sources = _finalize_overlap(result, counter, length_init, border)
 
     return _sources_to_dict(config, estimated_sources, source_indices)
 
 
-def demix_track_demucs(config, model, mix, device, pbar=False, source_indices=None):
+def demix_track_demucs(config, model, mix, device, pbar=False, source_indices=None, progress_callback=None):
     if _can_demix_mlx_full(model, device):
-        return demix_track_mlx_full(config, model, mix.cpu().numpy(), device, pbar=pbar, source_indices=source_indices)
+        return demix_track_mlx_full(config, model, mix.cpu().numpy(), device, pbar=pbar, source_indices=source_indices, progress_callback=progress_callback)
 
     source_indices = _normalize_source_indices(config, source_indices)
     source_names = _source_names(config)
@@ -543,7 +622,7 @@ def demix_track_demucs(config, model, mix, device, pbar=False, source_indices=No
             i = 0
             batch_data = []
             batch_locations = []
-            progress_bar = tqdm(total=mix.shape[1], desc="Processing audio chunks", leave=False) if pbar else None
+            progress = _ProgressContext(pbar, mix.shape[1], progress_callback)
 
             while i < mix.shape[1]:
                 part = mix[:, i:i + C].to(device)
@@ -563,11 +642,12 @@ def demix_track_demucs(config, model, mix, device, pbar=False, source_indices=No
                         counter[..., start:start+l] += 1.
                     batch_data, batch_locations = [], []
 
-                if progress_bar:
-                    progress_bar.update(step)
+                if progress.bar:
+                    progress.bar.update(step)
+                progress.emit(min(i, mix.shape[1]))
 
-            if progress_bar:
-                progress_bar.close()
+            progress.close()
+            progress.emit(mix.shape[1])
 
             estimated_sources = (result / counter).cpu().numpy()
             np.nan_to_num(estimated_sources, copy=False, nan=0.0)
@@ -576,13 +656,15 @@ def demix_track_demucs(config, model, mix, device, pbar=False, source_indices=No
         return estimated_sources
     return _sources_to_dict(config, estimated_sources, source_indices)
 
-def demix(config, model, mix: NDArray, device, pbar=False, model_type: str = None, source_indices=None) -> Dict[str, NDArray]:
+def demix(config, model, mix: NDArray, device, pbar=False, model_type: str = None, source_indices=None, progress_callback=None) -> Dict[str, NDArray]:
     if _can_demix_mlx_full(model, device):
-        return demix_track_mlx_full(config, model, mix, device, pbar=pbar, source_indices=source_indices)
+        return demix_track_mlx_full(config, model, mix, device, pbar=pbar, source_indices=source_indices, progress_callback=progress_callback)
     mix = torch.tensor(mix, dtype=torch.float32)
     if model_type in {'demucs', 'tasnet', 'legacy_demucs', 'legacy_tasnet'}:
         from .modules.legacy_demucs import apply_legacy_model
 
+        progress = _ProgressContext(callback=progress_callback)
+        progress.emit(0)
         with _autocast(device, config.training.get('use_amp', True)):
             with torch.inference_mode():
                 estimates = apply_legacy_model(
@@ -593,7 +675,8 @@ def demix(config, model, mix: NDArray, device, pbar=False, model_type: str = Non
                     overlap=float(config.inference.get('overlap', 0.25)),
                     progress=pbar,
                 ).cpu().numpy()
+        progress.emit(1)
         return dict(zip(config.training.instruments, estimates))
     if model_type == 'htdemucs':
-        return demix_track_demucs(config, model, mix, device, pbar=pbar, source_indices=source_indices)
-    return demix_track(config, model, mix, device, pbar=pbar, source_indices=source_indices)
+        return demix_track_demucs(config, model, mix, device, pbar=pbar, source_indices=source_indices, progress_callback=progress_callback)
+    return demix_track(config, model, mix, device, pbar=pbar, source_indices=source_indices, progress_callback=progress_callback)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pymss"
-version = "2.0.4"
+version = "2.0.5"
 description = "Python package for music source separation."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"