feat(inference): 新增 MPS MLX 推理缓存清理配置选项

- 【参数】注册 mps_mlx_clear_cache 到推理参数及透传列表
- 【校验】允许非 VR 模型类型使用该参数
- 【优化】新增上下文管理器以在评估后清理 MLX 缓存

Author: baicai-1145

pymss/separator.py

@@ -35,6 +35,7 @@
     "cuda_attention_backend": "inference",
     "mps_attention_backend": "inference",
     "mps_mlx_min_tokens": "inference",
+    "mps_mlx_clear_cache": "inference",
     "mps_model_backend": "inference",
     "mps_model_compute_dtype": "inference",
     "fuse_conv_bn": "inference",
@@ -55,6 +56,7 @@
         "use_amp",
         "cuda_attention_backend",
         "mps_attention_backend",
+        "mps_mlx_clear_cache",
         "mps_model_backend",
         "mps_model_compute_dtype",
         "fuse_conv_bn",

pymss/server/state.py

@@ -183,6 +183,8 @@ def _is_parameter_supported(config, model_type, key, section_name):
         bool: True when the condition is satisfied."""
     if model_type == "vr" and key in VR_SUPPORTED_PARAMETERS:
         return True
+    if key == "mps_mlx_clear_cache" and model_type != "vr":
+        return True
     # standardize is legacy input standardization backed by MSS YAML inference.normalize.
     # normalize is output peak normalization owned by runtime inference params.
     config_key = "normalize" if key == "standardize" else key

pymss/utils.py

@@ -792,7 +792,29 @@ def _mlx_fit_length(x, length):
     return x
 
 
-def _mlx_run_model_chunk(model, arr, chunk_size, progress_fraction_callback=None):
+@contextmanager
+def _mlx_clear_cache_after_eval(enabled=False):
+    """Clear MLX allocator cache after explicit eval points when requested."""
+    if not enabled:
+        yield
+        return
+    import mlx.core as mx
+
+    original_eval = mx.eval
+
+    def eval_and_clear(*args, **kwargs):
+        result = original_eval(*args, **kwargs)
+        clear_mlx_cache()
+        return result
+
+    mx.eval = eval_and_clear
+    try:
+        yield
+    finally:
+        mx.eval = original_eval
+
+
+def _mlx_run_model_chunk(model, arr, chunk_size, progress_fraction_callback=None, clear_cache_after_eval=False):
     """Implement the mlx run model chunk helper.
 
     Args:
@@ -802,7 +824,7 @@ def _mlx_run_model_chunk(model, arr, chunk_size, progress_fraction_callback=None
 
     Returns:
         Any: Computed result."""
-    with _model_progress_fraction_context(model, progress_fraction_callback):
+    with _mlx_clear_cache_after_eval(clear_cache_after_eval), _model_progress_fraction_context(model, progress_fraction_callback):
         y = model.mlx_forward_mx(arr)
     if y.ndim == arr.ndim:
         y = y[:, None]
@@ -926,6 +948,7 @@ def demix_track_mlx_full(config, model, mix, device, pbar=False, source_indices=
             mx.stack([chunk for (chunk, _), _ in batch], axis=0),
             C,
             lambda fraction: progress.emit(batch_done_before + round(batch_units * fraction)),
+            clear_cache_after_eval=bool(config.inference.get("mps_mlx_clear_cache", False)),
         )
         chunks = _mlx_select_sources(chunks, source_indices)
         for j, ((_, length), idx) in enumerate(batch):