Add DirectML support for PyMSS, support interrupting separation, fix CUDA Graph errors in input/output denoising, and update pip installation dependencies

README.md

@@ -191,6 +191,28 @@ hf download lj1995/VoiceConversionWebUI rmvpe.onnx --revision main \
   --local-dir assets/rmvpe
 ```
 
+### PyMSS DirectML 精度
+
+PyMSS 的 DirectML 模型参数精度默认为 `auto`：BS-Roformer 与
+Mel-Band-Roformer 会先在独立子进程中尝试原生 FP16 权重；若当前显卡、
+驱动或算子路径不兼容，当前子进程会先退出，再用全新的 FP32 子进程重试。
+该设置只影响 PyMSS 分离，不会改变 RVC 推理的全局精度。
+
+可在启动 WebUI 前设置 `PYMSS_DML_MODEL_DTYPE`：
+
+```bat
+rem 默认：FP16 优先，不兼容时回退 FP32
+set PYMSS_DML_MODEL_DTYPE=auto
+
+rem 始终使用 FP32，兼容性优先
+set PYMSS_DML_MODEL_DTYPE=float32
+
+rem 始终使用 FP16，不执行自动回退
+set PYMSS_DML_MODEL_DTYPE=float16
+```
+
+FP16 显存不足时不会自动改用占用更高的 FP32；界面会直接报告该错误。
+
 
 ### FFmpeg
 

i18n/locale/en_US.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…Omitted the first %s lines; showing the latest status only",
     "一键训练": "One-click training",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "Multiple audio files can also be imported. If a folder path exists, this input is ignored.",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "Batch vocal/accompaniment separation using UVR5 models.<br>You can choose a vocal-preserving model, or use DeEcho/DeReverb models to remove echo and reverb.",
     "仅支持pm和rmvpe音高提取算法": "Only the pm and rmvpe pitch extraction methods are supported",
     "从训练检查点提取的模型": "Model extracted from a training checkpoint",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "Enter the GPU index(es) separated by '-', e.g., 0-1-2 to use GPU 0, 1, and 2:",
     "任务": "Task",
-    "伴奏人声分离&去混响&去回声": "Vocals/Accompaniment Separation & Reverberation Removal",
     "使用显卡：%s": "GPUs in use: %s",
     "使用模型采样率": "Use model sample rate",
     "使用设备采样率": "Use device sample rate",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "Warming up CUDA Graph",
     "CUDA Graph预热完成": "CUDA Graph warm-up complete",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Data extraction started: start_time=%.6f, requested concurrency=%s, actual concurrency limit=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Data extraction finished: end_time=%.6f, total elapsed=%.3f seconds"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Data extraction finished: end_time=%.6f, total elapsed=%.3f seconds",
+    "人声伴奏分离&去混响": "Vocals/Accompaniment Separation & Dereverberation",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "Batch processing of vocals, accompaniment, and reverb using pymss/MSST models.",
+    "处理方式": "Processing mode",
+    "底层模型": "Underlying model",
+    "主结果文件夹": "Primary output folder",
+    "分离残余文件夹": "Residual output folder",
+    "停止分离": "Stop separation"
 }

i18n/locale/es_ES.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…Se omitieron las primeras %s líneas; se muestra solo el estado más reciente",
     "一键训练": "Entrenamiento con un clic",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "También se pueden importar varios archivos de audio. Si existe una ruta de carpeta, esta entrada se ignora.",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "Procesamiento por lotes para separar voces y acompañamiento mediante modelos UVR5.<br>Puede elegir un modelo que conserve las voces o usar modelos DeEcho/DeReverb para eliminar el eco y la reverberación.",
     "仅支持pm和rmvpe音高提取算法": "Solo se admiten los métodos de extracción de tono pm y rmvpe",
     "从训练检查点提取的模型": "Modeloo extraído de un punto de control de entrenamiento",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "Separe los números de identificación de la GPU con '-' al ingresarlos. Por ejemplo, '0-1-2' significa usar GPU 0, GPU 1 y GPU 2.",
     "任务": "Tarea",
-    "伴奏人声分离&去混响&去回声": "Separación de voz acompañante & eliminación de reverberación & eco",
     "使用显卡：%s": "GPU en uso: %s",
     "使用模型采样率": "Usar la frecuencia del modelo",
     "使用设备采样率": "Usar la frecuencia del dispositivo",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "Preparando CUDA Graph",
     "CUDA Graph预热完成": "Preparación de CUDA Graph completada",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Extracción de datos iniciada: start_time=%.6f, concurrencia solicitada=%s, límite real=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extracción de datos finalizada: end_time=%.6f, tiempo total=%.3f segundos"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extracción de datos finalizada: end_time=%.6f, tiempo total=%.3f segundos",
+    "人声伴奏分离&去混响": "Separación de voz/acompañamiento y eliminación de reverberación",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "Procesamiento por lotes de voz, acompañamiento y reverberación con modelos pymss/MSST.",
+    "处理方式": "Modo de procesamiento",
+    "底层模型": "Modelo subyacente",
+    "主结果文件夹": "Carpeta de salida principal",
+    "分离残余文件夹": "Carpeta de salida residual",
+    "停止分离": "Detener separación"
 }

i18n/locale/fr_FR.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…Les %s premières lignes ont été omises ; seul l'état récent est affiché",
     "一键训练": "Entraînement en un clic",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "Il est également possible d'importer plusieurs fichiers audio. Si un chemin de dossier existe, cette entrée est ignorée.",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "Séparation par lot des voix et de l’accompagnement à l’aide de modèles UVR5.<br>Vous pouvez choisir un modèle qui conserve les voix, ou utiliser les modèles DeEcho/DeReverb pour supprimer l’écho et la réverbération.",
     "仅支持pm和rmvpe音高提取算法": "Seules les méthodes d'extraction de hauteur pm et rmvpe sont prises en charge",
     "从训练检查点提取的模型": "Modèle extrait d'un point de contrôle d'entraînement",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "Entrez le(s) index GPU séparé(s) par '-', par exemple, 0-1-2 pour utiliser les GPU 0, 1 et 2 :",
     "任务": "Tâche",
-    "伴奏人声分离&去混响&去回声": "Séparation des voix/accompagnement et suppression de la réverbération",
     "使用显卡：%s": "GPU utilisés : %s",
     "使用模型采样率": "Utiliser la fréquence du modèle",
     "使用设备采样率": "Utiliser la fréquence du périphérique",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "Préchauffage de CUDA Graph",
     "CUDA Graph预热完成": "Préchauffage de CUDA Graph terminé",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Extraction des données démarrée : start_time=%.6f, parallélisme demandé=%s, limite réelle=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extraction des données terminée : end_time=%.6f, durée totale=%.3f secondes"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extraction des données terminée : end_time=%.6f, durée totale=%.3f secondes",
+    "人声伴奏分离&去混响": "Séparation voix/accompagnement et suppression de la réverbération",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "Traitement par lots des voix, de l’accompagnement et de la réverbération avec les modèles pymss/MSST.",
+    "处理方式": "Mode de traitement",
+    "底层模型": "Modèle sous-jacent",
+    "主结果文件夹": "Dossier de sortie principal",
+    "分离残余文件夹": "Dossier de sortie résiduel",
+    "停止分离": "Arrêter la séparation"
 }

i18n/locale/it_IT.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…Omesse le prime %s righe; viene mostrato solo lo stato più recente",
     "一键训练": "Addestramento con un clic",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "È anche possibile scegliere più file audio; la cartella ha la priorità",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "Separazione in batch di voce e accompagnamento tramite modelli UVR5.<br>È possibile scegliere un modello che preserva la voce oppure usare i modelli DeEcho/DeReverb per rimuovere eco e riverbero.",
     "仅支持pm和rmvpe音高提取算法": "Sono supportati solo i metodi di estrazione dell'intonazione pm e rmvpe",
     "从训练检查点提取的模型": "Modellolo estratto da un checkpoint di addestramento",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "Inserisci gli indici GPU separati da '-', ad esempio 0-1-2 per utilizzare GPU 0, 1 e 2:",
     "任务": "Attività",
-    "伴奏人声分离&去混响&去回声": "Separazione voce/accompagnamento",
     "使用显卡：%s": "GPU in uso: %s",
     "使用模型采样率": "Usa frequenza del modello",
     "使用设备采样率": "Usa frequenza del dispositivo",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "Riscaldamento di CUDA Graph",
     "CUDA Graph预热完成": "Riscaldamento di CUDA Graph completato",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Estrazione dati avviata: start_time=%.6f, concorrenza richiesta=%s, limite effettivo=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Estrazione dati completata: end_time=%.6f, tempo totale=%.3f secondi"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Estrazione dati completata: end_time=%.6f, tempo totale=%.3f secondi",
+    "人声伴奏分离&去混响": "Separazione voce/accompagnamento e rimozione del riverbero",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "Elaborazione batch di voce, accompagnamento e riverbero con i modelli pymss/MSST.",
+    "处理方式": "Modalità di elaborazione",
+    "底层模型": "Modello sottostante",
+    "主结果文件夹": "Cartella di output principale",
+    "分离残余文件夹": "Cartella di output residuo",
+    "停止分离": "Interrompi separazione"
 }

i18n/locale/ja_JP.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…先頭%s行を省略し、最新状態のみ表示しています",
     "一键训练": "ワンクリックトレーニング",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "複数のオーディオファイルをインポートすることもできます。フォルダパスが存在する場合、この入力は無視されます。",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "UVR5モデルを使用して、ボーカルと伴奏を一括分離します。<br>ボーカル保持モデルを選択するか、DeEcho／DeReverbモデルでエコーや残響を除去できます。",
     "仅支持pm和rmvpe音高提取算法": "ピッチ抽出方式はpmとrmvpeのみ対応しています",
     "从训练检查点提取的模型": "学習チェックポイントから抽出したモデル",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "ハイフンで区切って使用するGPUの番号を入力します。例えば0-1-2はGPU0、GPU1、GPU2を使用します",
     "任务": "タスク",
-    "伴奏人声分离&去混响&去回声": "伴奏ボーカル分離＆残響除去＆エコー除去",
     "使用显卡：%s": "使用GPU: %s",
     "使用模型采样率": "モデルのサンプルレートを使用",
     "使用设备采样率": "デバイスのサンプルレートを使用",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "CUDA Graphをウォームアップしています",
     "CUDA Graph预热完成": "CUDA Graphのウォームアップが完了しました",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "データ抽出開始：start_time=%.6f、要求並列数=%s、実並列上限=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "データ抽出完了：end_time=%.6f、合計所要時間=%.3f秒"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "データ抽出完了：end_time=%.6f、合計所要時間=%.3f秒",
+    "人声伴奏分离&去混响": "ボーカル・伴奏分離と残響除去",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "pymss/MSSTモデルを使用して、ボーカル、伴奏、残響を一括処理します。",
+    "处理方式": "処理方法",
+    "底层模型": "基盤モデル",
+    "主结果文件夹": "主要出力フォルダー",
+    "分离残余文件夹": "残余出力フォルダー",
+    "停止分离": "分離を停止"
 }

i18n/locale/ko_KR.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…앞의 %s줄을 생략하고 최신 상태만 표시합니다",
     "一键训练": "원클릭 훈련",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "여러 오디오 파일을 일괄 입력할 수도 있음, 둘 중 하나 선택, 폴더 우선 읽기",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "UVR5 모델을 사용하여 보컬과 반주를 일괄 분리합니다.<br>보컬 보존 모델을 선택하거나 DeEcho/DeReverb 모델로 에코와 잔향을 제거할 수 있습니다.",
     "仅支持pm和rmvpe音高提取算法": "피치 추출 방식은 pm과 rmvpe만 지원합니다",
     "从训练检查点提取的模型": "학습 체크포인트에서 추출한 모델",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "-로 구분하여 입력하는 카드 번호, 예: 0-1-2는 카드 0, 카드 1, 카드 2 사용",
     "任务": "작업",
-    "伴奏人声分离&去混响&去回声": "반주 인간 목소리 분리 & 혼효음 제거 & 에코 제거",
     "使用显卡：%s": "사용 GPU: %s",
     "使用模型采样率": "모델 샘플링 레이트 사용",
     "使用设备采样率": "장치 샘플링 레이트 사용",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "CUDA Graph 워밍업 중",
     "CUDA Graph预热完成": "CUDA Graph 워밍업 완료",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "데이터 추출 시작: start_time=%.6f, 요청 병렬 수=%s, 실제 병렬 상한=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "데이터 추출 완료: end_time=%.6f, 총 소요 시간=%.3f초"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "데이터 추출 완료: end_time=%.6f, 총 소요 시간=%.3f초",
+    "人声伴奏分离&去混响": "보컬/반주 분리 및 잔향 제거",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "pymss/MSST 모델로 보컬, 반주 및 잔향을 일괄 처리합니다.",
+    "处理方式": "처리 방식",
+    "底层模型": "기반 모델",
+    "主结果文件夹": "주 결과 폴더",
+    "分离残余文件夹": "잔여 결과 폴더",
+    "停止分离": "분리 중지"
 }

i18n/locale/pt_BR.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…As primeiras %s linhas foram omitidas; mostrando apenas o estado mais recente",
     "一键训练": "Treinamento com um clique",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "Você também pode inserir arquivos de áudio em lotes. Escolha uma das duas opções. É dada prioridade à leitura da pasta.",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "Processamento em lote para separar vocais e acompanhamento usando modelos UVR5.<br>Você pode escolher um modelo que preserve os vocais ou usar modelos DeEcho/DeReverb para remover eco e reverberação.",
     "仅支持pm和rmvpe音高提取算法": "Somente os métodos de extração de tom pm e rmvpe são compatíveis",
     "从训练检查点提取的模型": "Modeloo extraído de um checkpoint de treinamento",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "Digite o (s) índice(s) da GPU separados por '-', por exemplo, 0-1-2 para usar a GPU 0, 1 e 2:",
     "任务": "Tarefa",
-    "伴奏人声分离&去混响&去回声": "UVR5",
     "使用显卡：%s": "GPUs em uso: %s",
     "使用模型采样率": "Usar taxa do modelo",
     "使用设备采样率": "Usar taxa do dispositivo",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "Aquecendo o CUDA Graph",
     "CUDA Graph预热完成": "Aquecimento do CUDA Graph concluído",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Extração de dados iniciada: start_time=%.6f, concorrência solicitada=%s, limite real=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extração de dados concluída: end_time=%.6f, tempo total=%.3f segundos"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Extração de dados concluída: end_time=%.6f, tempo total=%.3f segundos",
+    "人声伴奏分离&去混响": "Separação de voz/acompanhamento e remoção de reverberação",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "Processamento em lote de voz, acompanhamento e reverberação com modelos pymss/MSST.",
+    "处理方式": "Modo de processamento",
+    "底层模型": "Modelo subjacente",
+    "主结果文件夹": "Pasta de saída principal",
+    "分离残余文件夹": "Pasta de saída residual",
+    "停止分离": "Parar separação"
 }

i18n/locale/ru_RU.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…Первые %s строк пропущены; показано только последнее состояние",
     "一键训练": "Обучение в одно нажатие",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "Можно также импортировать несколько аудиофайлов. Если путь к папке существует, то этот ввод игнорируется.",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "Пакетное разделение вокала и аккомпанемента с помощью моделей UVR5.<br>Можно выбрать модель с сохранением вокала или использовать модели DeEcho/DeReverb для удаления эха и реверберации.",
     "仅支持pm和rmvpe音高提取算法": "Поддерживаются только методы извлечения высоты тона pm и rmvpe",
     "从训练检查点提取的模型": "Модель извлечена из контрольной точки обучения",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "Введите, какие(-ую) GPU(-у) хотите использовать через '-', например 0-1-2, чтобы использовать GPU с номерами 0, 1 и 2:",
     "任务": "Задача",
-    "伴奏人声分离&去混响&去回声": "Разделение вокала/аккомпанемента и удаление эхо",
     "使用显卡：%s": "Используемые GPU: %s",
     "使用模型采样率": "Использовать частоту модели",
     "使用设备采样率": "Использовать частоту устройства",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "Прогрев CUDA Graph",
     "CUDA Graph预热完成": "Прогрев CUDA Graph завершён",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Извлечение данных начато: start_time=%.6f, запрошенный параллелизм=%s, фактический предел=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Извлечение данных завершено: end_time=%.6f, общее время=%.3f с"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Извлечение данных завершено: end_time=%.6f, общее время=%.3f с",
+    "人声伴奏分离&去混响": "Разделение вокала/аккомпанемента и удаление реверберации",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "Пакетная обработка вокала, аккомпанемента и реверберации с помощью моделей pymss/MSST.",
+    "处理方式": "Режим обработки",
+    "底层模型": "Базовая модель",
+    "主结果文件夹": "Основная папка вывода",
+    "分离残余文件夹": "Папка остаточного вывода",
+    "停止分离": "Остановить разделение"
 }

i18n/locale/tr_TR.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "…İlk %s satır atlandı; yalnızca en son durum gösteriliyor",
     "一键训练": "Tek Tuşla Eğit",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "Ses dosyaları ayrıca toplu olarak, iki seçimle, öncelikli okuma klasörüyle içe aktarılabilir",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "UVR5 modelleriyle vokal ve eşliği toplu olarak ayırır.<br>Vokali koruyan bir model seçebilir veya yankı ve reverbi kaldırmak için DeEcho/DeReverb modellerini kullanabilirsiniz.",
     "仅支持pm和rmvpe音高提取算法": "Yalnızca pm ve rmvpe perde çıkarma yöntemleri desteklenir",
     "从训练检查点提取的模型": "Eğitim kontrol noktasından çıkarılan model",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "GPU indekslerini '-' ile ayırarak girin, örneğin 0-1-2, GPU 0, 1 ve 2'yi kullanmak için:",
     "任务": "Görev",
-    "伴奏人声分离&去混响&去回声": "Vokal/Müzik Ayrıştırma ve Yankı Giderme",
     "使用显卡：%s": "Kullanılan GPU: %s",
     "使用模型采样率": "Model örnekleme hızını kullan",
     "使用设备采样率": "Cihaz örnekleme hızını kullan",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "CUDA Graph ısınıyor",
     "CUDA Graph预热完成": "CUDA Graph ısınması tamamlandı",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "Veri çıkarma başladı: start_time=%.6f, istenen eşzamanlılık=%s, gerçek üst sınır=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Veri çıkarma tamamlandı: end_time=%.6f, toplam süre=%.3f saniye"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "Veri çıkarma tamamlandı: end_time=%.6f, toplam süre=%.3f saniye",
+    "人声伴奏分离&去混响": "Vokal/altyapı ayırma ve yankı giderme",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "pymss/MSST modelleriyle vokal, altyapı ve yankıyı toplu olarak işler.",
+    "处理方式": "İşleme modu",
+    "底层模型": "Temel model",
+    "主结果文件夹": "Ana çıktı klasörü",
+    "分离残余文件夹": "Artık çıktı klasörü",
+    "停止分离": "Ayırmayı durdur"
 }

i18n/locale/zh_CN.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "……已省略前%s行，仅显示最新状态",
     "一键训练": "一键训练",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "也可批量输入音频文件, 二选一, 优先读文件夹",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。",
     "仅支持pm和rmvpe音高提取算法": "仅支持pm和rmvpe音高提取算法",
     "从训练检查点提取的模型": "从训练检查点提取的模型",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2",
     "任务": "任务",
-    "伴奏人声分离&去混响&去回声": "伴奏人声分离&去混响&去回声",
     "使用显卡：%s": "使用显卡：%s",
     "使用模型采样率": "使用模型采样率",
     "使用设备采样率": "使用设备采样率",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "正在预热CUDA Graph",
     "CUDA Graph预热完成": "CUDA Graph预热完成",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "数据提取结束：end_time=%.6f，总耗时=%.3f秒"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "数据提取结束：end_time=%.6f，总耗时=%.3f秒",
+    "人声伴奏分离&去混响": "人声伴奏分离&去混响",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "人声、伴奏与混响批量处理，使用pymss/MSST模型。",
+    "处理方式": "处理方式",
+    "底层模型": "底层模型",
+    "主结果文件夹": "主结果文件夹",
+    "分离残余文件夹": "分离残余文件夹",
+    "停止分离": "停止分离"
 }

i18n/locale/zh_HK.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "……已省略前%s行，僅顯示最新状态",
     "一键训练": "一鍵訓練",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "也可批量输入音频文件, 二选一, 优先读文件夹",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "批次處理人聲與伴奏分離，使用UVR5模型。<br>可選擇保留人聲的模型，或使用DeEcho、DeReverb模型去除回音和混響。",
     "仅支持pm和rmvpe音高提取算法": "僅支援pm和rmvpe音高提取算法",
     "从训练检查点提取的模型": "从訓練检查点提取的模型",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "以-分隔輸入使用的卡號, 例如 0-1-2 使用卡0和卡1和卡2",
     "任务": "任务",
-    "伴奏人声分离&去混响&去回声": "伴奏人聲分離&去混響&去回聲",
     "使用显卡：%s": "使用顯示卡：%s",
     "使用模型采样率": "使用模型采样率",
     "使用设备采样率": "使用设备采样率",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "正在預熱 CUDA Graph",
     "CUDA Graph预热完成": "CUDA Graph 預熱完成",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "資料擷取開始：start_time=%.6f，請求並行數=%s，實際並行數上限=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "資料擷取結束：end_time=%.6f，總耗時=%.3f秒"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "資料擷取結束：end_time=%.6f，總耗時=%.3f秒",
+    "人声伴奏分离&去混响": "人聲伴奏分離&去混響",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "使用 pymss/MSST 模型批次處理人聲、伴奏與混響。",
+    "处理方式": "處理方式",
+    "底层模型": "底層模型",
+    "主结果文件夹": "主要結果資料夾",
+    "分离残余文件夹": "分離殘餘資料夾",
+    "停止分离": "停止分離"
 }

i18n/locale/zh_SG.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "……已省略前%s行，仅显示最新状态",
     "一键训练": "一鍵訓練",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "也可批量输入音频文件, 二选一, 优先读文件夹",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。",
     "仅支持pm和rmvpe音高提取算法": "仅支持pm和rmvpe音高提取算法",
     "从训练检查点提取的模型": "从训练检查点提取的模型",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "以-分隔輸入使用的卡號, 例如 0-1-2 使用卡0和卡1和卡2",
     "任务": "任务",
-    "伴奏人声分离&去混响&去回声": "伴奏人聲分離&去混響&去回聲",
     "使用显卡：%s": "使用显卡：%s",
     "使用模型采样率": "使用模型采样率",
     "使用设备采样率": "使用设备采样率",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "正在预热CUDA Graph",
     "CUDA Graph预热完成": "CUDA Graph预热完成",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "数据提取结束：end_time=%.6f，总耗时=%.3f秒"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "数据提取结束：end_time=%.6f，总耗时=%.3f秒",
+    "人声伴奏分离&去混响": "人声伴奏分离&去混响",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "人声、伴奏与混响批量处理，使用pymss/MSST模型。",
+    "处理方式": "处理方式",
+    "底层模型": "底层模型",
+    "主结果文件夹": "主结果文件夹",
+    "分离残余文件夹": "分离残余文件夹",
+    "停止分离": "停止分离"
 }

i18n/locale/zh_TW.json

@@ -64,12 +64,10 @@
     "……已省略前%s行，仅显示最新状态": "……已省略前%s行，僅顯示最新状态",
     "一键训练": "一鍵訓練",
     "也可批量输入音频文件, 二选一, 优先读文件夹": "也可批量输入音频文件, 二选一, 优先读文件夹",
-    "人声伴奏分离批量处理，使用UVR5模型。<br>可选择保留人声模型，或使用DeEcho、DeReverb模型去除延迟和混响。": "批次處理人聲與伴奏分離，使用UVR5模型。<br>可選擇保留人聲的模型，或使用DeEcho、DeReverb模型去除回音和混響。",
     "仅支持pm和rmvpe音高提取算法": "僅支援pm和rmvpe音高提取算法",
     "从训练检查点提取的模型": "从訓練检查点提取的模型",
     "以-分隔输入使用的卡号, 例如   0-1-2   使用卡0和卡1和卡2": "以-分隔輸入使用的卡號, 例如 0-1-2 使用卡0和卡1和卡2",
     "任务": "任务",
-    "伴奏人声分离&去混响&去回声": "伴奏人聲分離&去混響&去回聲",
     "使用显卡：%s": "使用顯示卡：%s",
     "使用模型采样率": "使用模型采样率",
     "使用设备采样率": "使用设备采样率",
@@ -261,5 +259,12 @@
     "正在预热CUDA Graph": "正在預熱 CUDA Graph",
     "CUDA Graph预热完成": "CUDA Graph 預熱完成",
     "数据提取开始：start_time=%.6f，请求并行数=%s，实际并行数上限=%s": "資料擷取開始：start_time=%.6f，請求並行數=%s，實際並行數上限=%s",
-    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "資料擷取結束：end_time=%.6f，總耗時=%.3f秒"
+    "数据提取结束：end_time=%.6f，总耗时=%.3f秒": "資料擷取結束：end_time=%.6f，總耗時=%.3f秒",
+    "人声伴奏分离&去混响": "人聲伴奏分離&去混響",
+    "人声、伴奏与混响批量处理，使用pymss/MSST模型。": "使用 pymss/MSST 模型批次處理人聲、伴奏與混響。",
+    "处理方式": "處理方式",
+    "底层模型": "底層模型",
+    "主结果文件夹": "主要結果資料夾",
+    "分离残余文件夹": "分離殘餘資料夾",
+    "停止分离": "停止分離"
 }

realtime_gui.py

@@ -675,6 +675,8 @@ if __name__ == "__main__":
                 ).to(self.config.device)
             else:
                 self.resampler2 = None
+            # Bundled torch.istft is not CUDA Graph-capturable, so TorchGate
+            # stays eager while resampling and RVC inference still use graphs.
             self.tg = TorchGate(
                 sr=self.gui_config.samplerate, n_fft=4 * self.zc, prop_decrease=0.9
             ).to(self.config.device)
@@ -697,15 +699,7 @@ if __name__ == "__main__":
                     short = self.input_wav[
                         -self.sola_buffer_frame - self.block_frame :
                     ].unsqueeze(0)
-                    run_cuda_graph(
-                        self.tg,
-                        "realtime-input-noise-reduction",
-                        lambda short_audio, full_audio: self.tg(
-                            short_audio, full_audio
-                        ),
-                        short,
-                        self.input_wav.unsqueeze(0),
-                    )
+                    self.tg(short, self.input_wav.unsqueeze(0))
 
                 resample_input = self.input_wav[-self.block_frame - 2 * self.zc :]
                 run_cuda_graph(
@@ -730,15 +724,7 @@ if __name__ == "__main__":
                         inferred,
                     )
                 if self.gui_config.O_noise_reduce:
-                    run_cuda_graph(
-                        self.tg,
-                        "realtime-output-noise-reduction",
-                        lambda short_audio, full_audio: self.tg(
-                            short_audio, full_audio
-                        ),
-                        inferred.unsqueeze(0),
-                        self.output_buffer.unsqueeze(0),
-                    )
+                    self.tg(inferred.unsqueeze(0), self.output_buffer.unsqueeze(0))
                 torch.cuda.synchronize(self.config.device)
                 printt(i18n("CUDA Graph预热完成"))
             except Exception:
@@ -821,12 +807,8 @@ if __name__ == "__main__":
                     self.block_frame :
                 ].clone()
                 input_wav = self.input_wav[-self.sola_buffer_frame - self.block_frame :]
-                input_wav = run_cuda_graph(
-                    self.tg,
-                    "realtime-input-noise-reduction",
-                    lambda short, full: self.tg(short, full),
-                    input_wav.unsqueeze(0),
-                    self.input_wav.unsqueeze(0),
+                input_wav = self.tg(
+                    input_wav.unsqueeze(0), self.input_wav.unsqueeze(0)
                 ).squeeze(0)
                 input_wav[: self.sola_buffer_frame] *= self.fade_in_window
                 input_wav[: self.sola_buffer_frame] += (
@@ -879,12 +861,8 @@ if __name__ == "__main__":
                     self.block_frame :
                 ].clone()
                 self.output_buffer[-self.block_frame :] = infer_wav[-self.block_frame :]
-                infer_wav = run_cuda_graph(
-                    self.tg,
-                    "realtime-output-noise-reduction",
-                    lambda short, full: self.tg(short, full),
-                    infer_wav.unsqueeze(0),
-                    self.output_buffer.unsqueeze(0),
+                infer_wav = self.tg(
+                    infer_wav.unsqueeze(0), self.output_buffer.unsqueeze(0)
                 ).squeeze(0)
             # volume envelop mixing
             if self.gui_config.rms_mix_rate < 1 and self.function == "vc":

tools/process_utils.py

@@ -0,0 +1,71 @@
+import logging
+import os
+import signal
+import subprocess
+import time
+
+
+logger = logging.getLogger(__name__)
+
+
+def kill_process_tree(process, process_name="", task_logger=None):
+    """Terminate a Popen process and every child process it created."""
+    if process is None:
+        return False
+    try:
+        if process.poll() is not None:
+            return False
+    except (OSError, ProcessLookupError):
+        return False
+
+    pid = process.pid
+    if os.name == "nt":
+        try:
+            subprocess.run(
+                ["taskkill", "/t", "/f", "/pid", str(pid)],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                check=False,
+            )
+        except OSError:
+            try:
+                process.terminate()
+            except (OSError, ProcessLookupError):
+                pass
+    else:
+        try:
+            process_group = os.getpgid(pid)
+        except (OSError, ProcessLookupError):
+            process_group = None
+        try:
+            if process_group is not None and process_group != os.getpgrp():
+                os.killpg(process_group, signal.SIGTERM)
+            else:
+                os.kill(pid, signal.SIGTERM)
+        except (OSError, ProcessLookupError):
+            pass
+
+        for _ in range(10):
+            if process.poll() is not None:
+                break
+            time.sleep(0.1)
+        if process.poll() is None:
+            try:
+                if process_group is not None and process_group != os.getpgrp():
+                    os.killpg(process_group, signal.SIGKILL)
+                else:
+                    os.kill(pid, signal.SIGKILL)
+            except (OSError, ProcessLookupError):
+                pass
+
+    try:
+        process.wait(timeout=5)
+    except (OSError, ProcessLookupError, subprocess.TimeoutExpired):
+        try:
+            process.kill()
+        except (OSError, ProcessLookupError):
+            pass
+
+    log = task_logger or logger
+    log.info("%s process tree terminated (pid=%s)", process_name or "Child", pid)
+    return True

tools/pymss/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 KitsuneX07
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tools/pymss/__init__.py

@@ -0,0 +1,68 @@
+"""Public Python API for pymss.
+
+pymss provides model catalog helpers, model downloading, audio I/O, ensemble
+utilities, logging helpers, and the ``MSSeparator`` runtime for music source
+separation. Most users can import from this top-level package instead of
+importing submodules directly.
+
+Exports:
+    MSSeparator: Main runtime class for loading separation models and producing
+        stems. Prefer ``MSSeparator.from_model_name(...)`` for catalog models.
+    get_separation_logger: Create or reuse the package logger.
+    create_separator: Create ``MSSeparator`` from a catalog model name.
+    get_model_entry: Resolve catalog metadata for one model name or alias.
+    list_models: List model catalog entries.
+    resolve_model: Resolve catalog model paths without constructing a
+        separator.
+    download_model: Download all files required by one catalog model.
+    ensemble_audios: Load and combine multiple audio files.
+    save_ensemble_audio: Ensemble multiple audio files and save the result.
+    WorkflowRunner: Run a multi-model audio workflow.
+    load_audio: Load an audio file into a NumPy array.
+    save_audio: Save a NumPy audio array to wav/flac/mp3/m4a.
+
+Example:
+    >>> from pymss import MSSeparator
+    >>> separator = MSSeparator.from_model_name(
+    ...     "bs_roformer_voc_hyperacev2",
+    ...     download=True,
+    ...     model_dir="models",
+    ... )
+    >>> separator.process_folder("song.wav")
+
+Example:
+    >>> from pymss import download_model, list_models
+    >>> models = list_models(supported=True)
+    >>> download_model(models[0].name, model_dir="models")
+
+Example:
+    >>> from pymss import ensemble_audios, save_ensemble_audio
+    >>> audio, sample_rate = ensemble_audios(["a.wav", "b.wav"], weights=[1, 1])
+    >>> save_ensemble_audio(["a.wav", "b.wav"], "ensemble.wav")
+"""
+
+from .separator import MSSeparator
+from .logger import get_separation_logger
+from .model_registry import create_separator, get_model_entry, list_models, resolve_model
+from .model_download import download_model
+from .ensemble import ensemble_audios, save_ensemble_audio
+from .audio_io import load_audio, save_audio
+from .workflow import WorkflowRunner, load_workflow_file, run_workflow_file, validate_workflow
+
+__all__ = (
+    "MSSeparator",
+    "get_separation_logger",
+    "create_separator",
+    "get_model_entry",
+    "list_models",
+    "resolve_model",
+    "download_model",
+    "ensemble_audios",
+    "save_ensemble_audio",
+    "WorkflowRunner",
+    "load_workflow_file",
+    "run_workflow_file",
+    "validate_workflow",
+    "load_audio",
+    "save_audio",
+)

tools/pymss/audio_io.py

@@ -0,0 +1,280 @@
+import json
+import subprocess
+
+import av
+import numpy as np
+
+
+def _frame_to_audio(frame, mono):
+    """Implement the frame to audio helper.
+
+    Args:
+        frame (Any): Frame value.
+        mono (bool): Mono value.
+
+    Returns:
+        Any: Computed result."""
+    audio = frame.to_ndarray()
+    audio = audio[None, :] if audio.ndim == 1 else audio
+    return (audio.mean(axis=0, keepdims=True) if mono and audio.shape[0] > 1 else audio).astype(np.float32, copy=False)
+
+
+def _ffmpeg_audio_stream_info(path):
+    """Return basic audio stream information from ffprobe."""
+    command = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "a:0",
+        "-show_entries",
+        "stream=sample_rate,channels",
+        "-of",
+        "json",
+        str(path),
+    ]
+    result = subprocess.run(command, check=True, capture_output=True, text=True)
+    streams = json.loads(result.stdout or "{}").get("streams") or []
+    if not streams:
+        raise ValueError(f"No audio stream found in {path!s}.")
+    stream = streams[0]
+    return int(stream["sample_rate"]), int(stream["channels"])
+
+
+def _load_audio_ffmpeg(path, sr=None, mono=False, offset=0.0, duration=None):
+    """Load audio through the ffmpeg CLI as a fallback for damaged streams."""
+    source_rate, source_channels = _ffmpeg_audio_stream_info(path)
+    out_rate = int(sr or source_rate)
+    channels = 1 if mono else source_channels
+    command = ["ffmpeg", "-nostdin", "-v", "error"]
+    if offset:
+        command += ["-ss", str(float(offset))]
+    command += ["-i", str(path), "-map", "0:a:0", "-vn"]
+    if duration is not None:
+        command += ["-t", str(float(duration))]
+    command += ["-f", "f32le", "-acodec", "pcm_f32le", "-ar", str(out_rate), "-ac", str(channels), "-"]
+    result = subprocess.run(command, check=True, capture_output=True)
+    audio = np.frombuffer(result.stdout, dtype="<f4")
+    complete_samples = audio.size // channels
+    audio = audio[: complete_samples * channels]
+    audio = audio.reshape(complete_samples, channels).T
+    audio = np.ascontiguousarray(audio.astype(np.float32, copy=False))
+    return (audio[0] if mono or channels == 1 else audio), out_rate
+
+
+def _load_audio_librosa(path, sr=None, mono=False, offset=0.0, duration=None):
+    """Load audio through librosa as a fallback before the ffmpeg CLI."""
+    import librosa
+
+    audio, out_rate = librosa.load(
+        path,
+        sr=sr,
+        mono=mono,
+        offset=float(offset or 0.0),
+        duration=None if duration is None else float(duration),
+        dtype=np.float32,
+    )
+    audio = np.ascontiguousarray(np.asarray(audio, dtype=np.float32))
+    return (audio[0] if audio.ndim > 1 and (mono or audio.shape[0] == 1) else audio), int(out_rate)
+
+
+def _load_audio_av(path, sr=None, mono=False, offset=0.0, duration=None):
+    """Load audio through PyAV."""
+    chunks = []
+    out_rate = None
+    with av.open(path) as container:
+        stream = container.streams.audio[0]
+        out_rate = int(sr or stream.rate)
+        resampler = None
+        stop_samples = None if duration is None else int(round((offset + duration) * out_rate))
+        decoded = 0
+
+        for frame in container.decode(stream):
+            if resampler is None:
+                resampler = av.AudioResampler(format="fltp", layout=frame.layout.name, rate=out_rate)
+
+            for out in resampler.resample(frame):
+                chunks.append(audio := _frame_to_audio(out, mono))
+                decoded += audio.shape[-1]
+            if stop_samples is not None and decoded >= stop_samples:
+                break
+
+        if resampler is not None:
+            for out in resampler.resample(None):
+                chunks.append(_frame_to_audio(out, mono))
+
+    start = int(round(offset * out_rate))
+    stop = None if duration is None else start + int(round(duration * out_rate))
+    channels = 1 if mono else 0
+    audio = np.ascontiguousarray(
+        (np.concatenate(chunks, axis=-1) if chunks else np.empty((channels, 0), dtype=np.float32))[..., start:stop]
+    )
+    return (audio[0] if mono or audio.shape[0] == 1 else audio), out_rate
+
+
+def load_audio(path, sr=None, mono=False, offset=0.0, duration=None):
+    """Load an audio file as float32 NumPy samples.
+
+    Audio decoding is attempted in this order: PyAV, librosa, then the
+    ffmpeg CLI fallback. Stereo or multi-channel output is returned
+    channel-first as ``(channels, samples)``. Mono output is returned as a
+    one-dimensional array.
+
+    Args:
+        path (str | os.PathLike): Input audio file path. Any format supported
+            by the local FFmpeg/PyAV build can be decoded.
+        sr (int | None, optional): Target sample rate. ``None`` keeps the
+            source stream sample rate. Defaults to None.
+        mono (bool, optional): Whether to downmix multi-channel audio to mono.
+            Defaults to False.
+        offset (float, optional): Start offset in seconds. Defaults to 0.0.
+        duration (float | None, optional): Maximum duration to return in
+            seconds after ``offset``. ``None`` reads to the end. Defaults to
+            None.
+
+    Returns:
+        tuple[np.ndarray, int]: Audio samples and sample rate. The array is
+        channel-first for multi-channel audio and one-dimensional for mono
+        output.
+
+    Example:
+        >>> from pymss import load_audio
+        >>> audio, sample_rate = load_audio("song.wav", sr=44100)
+        >>> sample_rate
+        44100
+
+    Example:
+        >>> clip, sample_rate = load_audio(
+        ...     "song.wav",
+        ...     mono=True,
+        ...     offset=30.0,
+        ...     duration=10.0,
+        ... )
+        >>> clip.ndim
+        1"""
+    loaders = [
+        ("PyAV", _load_audio_av),
+        ("librosa", _load_audio_librosa),
+        ("ffmpeg CLI", _load_audio_ffmpeg),
+    ]
+
+    errors = []
+    for name, loader in loaders:
+        try:
+            return loader(path, sr=sr, mono=mono, offset=offset, duration=duration)
+        except Exception as e:
+            errors.append(f"{name}: {e}")
+            continue
+
+    raise RuntimeError(f"All audio loading methods failed: {'; '.join(errors)}")
+
+
+def _bitrate_to_int(value):
+    """Implement the bitrate to int helper.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Computed result."""
+    if value is None:
+        return None
+    if isinstance(value, int):
+        return value
+    value = str(value).strip().lower()
+    return int(float(value[:-1]) * 1000) if value.endswith("k") else int(value)
+
+
+def _format_audio(audio):
+    """Format audio.
+
+    Args:
+        audio (np.ndarray): Audio samples.
+
+    Returns:
+        Any: Computed result."""
+    audio = np.asarray(audio)
+    audio = np.ascontiguousarray(audio[:, None] if audio.ndim == 1 else audio)
+    # We can use "fltp" container for all output formats, while the final result is determined by the codec.
+    # Using the fltp sample format can also help avoid some clipping distortion that occurs with integer formats.
+    return np.ascontiguousarray(audio.astype(np.float32).T)
+
+
+def save_audio(path, audio, sr, output_format, audio_params):
+    """Save a NumPy audio array to wav, flac, mp3, or m4a.
+
+    Audio is expected as sample-major data, either ``(samples,)`` for mono or
+    ``(samples, channels)`` for multi-channel audio. The output codec is chosen
+    from ``output_format`` and ``audio_params``.
+
+    Args:
+        path (str | os.PathLike): Output file path.
+        audio (np.ndarray): Audio samples. Mono arrays may be one-dimensional;
+            stereo arrays should be shaped as ``(samples, 2)``.
+        sr (int): Sample rate in Hz.
+        output_format (str): Output format. Supported values are ``wav``,
+            ``flac``, ``mp3``, and ``m4a``.
+        audio_params (dict): Encoding options. Supported keys include
+            ``wav_bit_depth`` (``FLOAT``, ``PCM_16``, ``PCM_24``),
+            ``flac_bit_depth`` (currently ``PCM_24`` uses soundfile),
+            ``mp3_bit_rate`` (for example ``"320k"``), ``m4a_bit_rate``,
+            ``m4a_codec``, and ``m4a_aac_at_quality``.
+
+    Returns:
+        None: The file is written to ``path``.
+
+    Example:
+        >>> from pymss import save_audio
+        >>> save_audio(
+        ...     "vocals.wav",
+        ...     vocals,
+        ...     44100,
+        ...     "wav",
+        ...     {"wav_bit_depth": "FLOAT"},
+        ... )
+
+    Example:
+        >>> save_audio(
+        ...     "instrumental.flac",
+        ...     instrumental,
+        ...     44100,
+        ...     "flac",
+        ...     {"flac_bit_depth": "PCM_24"},
+        ... )"""
+    output_format = output_format.lower()
+    audio_array = np.asarray(audio)
+    layout = "stereo" if audio_array.ndim > 1 and audio_array.shape[1] == 2 else "mono"
+
+    if output_format == "mp3":
+        codec = "libmp3lame"
+    elif output_format == "m4a":
+        codec = audio_params.get("m4a_codec", "aac")
+    elif output_format == "flac":
+        # PyAV's FLAC encoder only exposes a single "flac" codec in the current version.
+        # In the current version, without access to bits_per_raw_sample in PyAV, PCM_24 may still be encoded as 16-bit.
+        # Use soundfile to save 24-bit FLAC
+        codec = "flac"
+        if audio_params.get("flac_bit_depth", "PCM_24") == "PCM_24":
+            import soundfile as sf
+
+            return sf.write(path, audio_array, int(sr), format="FLAC", subtype="PCM_24")
+    else:
+        wav_codecs = {"PCM_16": "pcm_s16le", "PCM_24": "pcm_s24le", "FLOAT": "pcm_f32le"}
+        codec = wav_codecs.get(audio_params.get("wav_bit_depth", "FLOAT"), wav_codecs["FLOAT"])
+
+    with av.open(path, "w") as container:
+        stream = container.add_stream(codec, rate=int(sr))
+        stream.layout = layout
+        if output_format == "mp3":
+            stream.bit_rate = _bitrate_to_int(audio_params.get("mp3_bit_rate", "320k"))
+        elif output_format == "m4a":
+            stream.bit_rate = _bitrate_to_int(audio_params.get("m4a_bit_rate", "512k"))
+            if codec == "aac_at":
+                stream.codec_context.options = {"aac_at_quality": str(audio_params.get("m4a_aac_at_quality", 2))}
+
+        frame = av.AudioFrame.from_ndarray(_format_audio(audio_array), format="fltp", layout=layout)
+        frame.sample_rate = int(sr)
+        for packet in stream.encode(frame):
+            container.mux(packet)
+        for packet in stream.encode():
+            container.mux(packet)

tools/pymss/cli.py

@@ -0,0 +1,587 @@
+import argparse
+import json
+import sys
+import warnings
+
+from .ensemble import ENSEMBLE_ALGORITHMS, save_ensemble_audio
+from .logger import get_separation_logger
+from .model_download import download_all, download_model
+from .model_registry import create_separator, list_models, resolve_model
+from .progress import _CliInferenceProgress
+from .workflow import load_workflow_file, run_workflow_file, validate_workflow, write_workflow_template
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+def _parse_key_value(values):
+    """Parse key value.
+
+    Args:
+        values (Any): Values value.
+
+    Returns:
+        Any: Parsed value."""
+    result = {}
+    for value in values or []:
+        if "=" not in value:
+            raise argparse.ArgumentTypeError(f"Expected key=value, got {value!r}")
+        key, raw = value.split("=", 1)
+        lowered = raw.lower()
+        if lowered in {"true", "false"}:
+            result[key] = lowered == "true"
+        else:
+            try:
+                result[key] = int(raw)
+            except ValueError:
+                try:
+                    result[key] = float(raw)
+                except ValueError:
+                    result[key] = raw
+    return result
+
+
+def cmd_list(args):
+    """Implement the cmd list helper.
+
+    Args:
+        args (argparse.Namespace): Parsed command-line arguments.
+
+    Returns:
+        Any: Computed result."""
+    rows = list_models(category=args.category, supported=None if args.all else True)
+    if args.json:
+        print(json.dumps([item.__dict__ for item in rows], ensure_ascii=False, indent=2))
+        return 0
+    for item in rows:
+        status = "ok" if item.supported else item.unsupported_reason
+        category = item.category_path or item.primary_category
+        print(f"{item.name}\t{item.model_type or item.architecture}\t{category}\t{item.target_stem}\t{status}")
+    return 0
+
+
+def cmd_info(args):
+    """Implement the cmd info helper.
+
+    Args:
+        args (argparse.Namespace): Parsed command-line arguments.
+
+    Returns:
+        Any: Computed result."""
+    resolved = resolve_model(args.model, model_dir=args.model_dir, require_supported=False, require_exists=False)
+    entry = resolved["entry"]
+    data = {
+        "name": entry.name,
+        "model_type": entry.model_type,
+        "architecture": entry.architecture,
+        "supported": entry.supported,
+        "unsupported_reason": entry.unsupported_reason,
+        "category": entry.category_path or entry.primary_category,
+        "category_cn": " / ".join(filter(None, [entry.primary_category_cn, entry.secondary_category_cn])),
+        "target_stem": entry.target_stem,
+        "model_path": resolved["model_path"],
+        "config_path": resolved["config_path"],
+        "size_bytes": entry.size_bytes,
+    }
+    print(json.dumps(data, ensure_ascii=False, indent=2))
+    return 0
+
+
+def cmd_download(args):
+    """Implement the cmd download helper.
+
+    Args:
+        args (argparse.Namespace): Parsed command-line arguments.
+
+    Returns:
+        Any: Computed result."""
+    if args.model == "all":
+        results = download_all(
+            model_dir=args.model_dir,
+            source=args.source,
+            endpoint=args.endpoint,
+            supported_only=args.supported_only,
+            force=args.force,
+        )
+        failed = [item for item in results if item.get("error")]
+        print(f"Downloaded/skipped {len(results) - len(failed)} model(s), failed {len(failed)}.")
+        for item in failed:
+            print(f"ERROR {item['entry'].name}: {item['error']}", file=sys.stderr)
+        return 1 if failed else 0
+
+    result = download_model(
+        args.model,
+        model_dir=args.model_dir,
+        source=args.source,
+        endpoint=args.endpoint,
+        force=args.force,
+    )
+    _print_download_result(result)
+    return 0
+
+
+def _print_download_result(result):
+    """Print download result.
+
+    Args:
+        result (Any): Result value.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    for path in result["skipped"]:
+        print(f"exists {path}")
+    for path in result["downloaded"]:
+        print(f"downloaded {path}")
+
+
+def _ensure_model_files(args):
+    """Ensure model files.
+
+    Args:
+        args (argparse.Namespace): Parsed command-line arguments.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    try:
+        resolve_model(args.model, model_dir=args.model_dir, require_supported=True, require_exists=True)
+    except FileNotFoundError:
+        result = download_model(args.model, model_dir=args.model_dir, source=args.source, endpoint=args.endpoint)
+        _print_download_result(result)
+    else:
+        if args.download:
+            result = download_model(args.model, model_dir=args.model_dir, source=args.source, endpoint=args.endpoint)
+            _print_download_result(result)
+
+
+def cmd_infer(args):
+    """Implement the cmd infer helper.
+
+    Args:
+        args (argparse.Namespace): Parsed command-line arguments.
+
+    Returns:
+        Any: Computed result."""
+    _ensure_model_files(args)
+    logger = get_separation_logger()
+    inference_progress = _CliInferenceProgress()
+    with create_separator(
+        args.model,
+        model_dir=args.model_dir,
+        device=args.device,
+        device_ids=args.device_ids or [0],
+        output_format=args.output_format,
+        audio_params={
+            "wav_bit_depth": args.wav_bit_depth,
+            "flac_bit_depth": args.flac_bit_depth,
+            "mp3_bit_rate": args.mp3_bit_rate,
+            "m4a_bit_rate": args.m4a_bit_rate,
+            "m4a_aac_at_quality": args.m4a_aac_at_quality,
+        },
+        use_tta=args.tta,
+        store_dirs=args.output,
+        save_as_folder=args.save_as_folder,
+        logger=logger,
+        debug=args.debug,
+        progress_callback=inference_progress,
+        inference_params=_parse_key_value(args.param),
+    ) as separator:
+        try:
+            files = separator.process_folder(args.input)
+        finally:
+            inference_progress.close()
+    logger.info(f"Processed {len(files)} file(s).")
+    return 0
+
+
+def cmd_ensemble(args):
+    """Run the ensemble CLI command.
+
+    Args:
+        args (argparse.Namespace): Parsed command-line arguments.
+
+    Returns:
+        Any: Computed result."""
+    logger = get_separation_logger()
+    output_path = save_ensemble_audio(
+        args.files,
+        args.output,
+        algorithm=args.algorithm,
+        weights=args.weights,
+        output_format=args.output_format,
+        audio_params={
+            "wav_bit_depth": args.wav_bit_depth,
+            "flac_bit_depth": args.flac_bit_depth,
+            "mp3_bit_rate": args.mp3_bit_rate,
+            "m4a_bit_rate": args.m4a_bit_rate,
+            "m4a_codec": args.m4a_codec,
+            "m4a_aac_at_quality": args.m4a_aac_at_quality,
+        },
+        logger=logger,
+    )
+    logger.info(f"Saved ensemble audio to {output_path}")
+    return 0
+
+
+def cmd_workflow_init(args):
+    """Write a starter workflow file."""
+    path = write_workflow_template(args.output, overwrite=args.force)
+    print(f"Wrote workflow template to {path}")
+    return 0
+
+
+def cmd_workflow_validate(args):
+    """Validate a workflow file without running inference."""
+    workflow = load_workflow_file(args.config)
+    model_resolver = resolve_model if args.check_models or args.require_files else None
+    validate_workflow(
+        workflow,
+        model_dir=args.model_dir,
+        require_model_files=args.require_files,
+        model_resolver=model_resolver,
+    )
+    print(f"Workflow is valid: {len(workflow.steps)} step(s).")
+    return 0
+
+
+def cmd_workflow_run(args):
+    """Run an audio workflow from a YAML/JSON file."""
+    logger = get_separation_logger()
+    files = run_workflow_file(
+        args.config,
+        args.input,
+        args.output,
+        model_dir=args.model_dir,
+        device=args.device,
+        output_format=args.output_format,
+        download=args.download,
+        source=args.source,
+        endpoint=args.endpoint,
+        output_layout=args.output_layout,
+        audio_params={
+            "wav_bit_depth": args.wav_bit_depth,
+            "flac_bit_depth": args.flac_bit_depth,
+            "mp3_bit_rate": args.mp3_bit_rate,
+            "m4a_bit_rate": args.m4a_bit_rate,
+            "m4a_codec": args.m4a_codec,
+            "m4a_aac_at_quality": args.m4a_aac_at_quality,
+        },
+        logger=logger,
+        debug=args.debug,
+    )
+    logger.info(f"Processed {len(files)} file(s).")
+    return 0
+
+
+def cmd_serve(args):
+    """Implement the cmd serve helper.
+
+    Args:
+        args (argparse.Namespace): Parsed command-line arguments.
+
+    Returns:
+        Any: Computed result."""
+    from .server import ServerConfig, run_server
+
+    config = ServerConfig(
+        model=args.model,
+        model_dir=args.model_dir,
+        source=args.source,
+        endpoint=args.endpoint,
+        device=args.device,
+        device_ids=args.device_ids or [0],
+        api_key=args.api_key,
+        host=args.host,
+        port=args.port,
+        debug=args.debug,
+        inference_params=_parse_key_value(args.param),
+        max_audio_seconds=args.max_audio_seconds,
+        max_request_bytes=args.max_request_bytes,
+        max_queue_size=args.max_queue_size,
+        request_timeout_seconds=args.request_timeout_seconds,
+        webui=args.webui,
+    )
+    run_server(config)
+    return 0
+
+
+def build_parser():
+    """Build the pymss command-line parser.
+
+    Args:
+        None: This callable does not accept user-provided arguments.
+
+    Returns:
+        argparse.ArgumentParser: Configured CLI parser."""
+    parser = argparse.ArgumentParser(
+        prog="pymss",
+        description="Command-line interface for the pymss music source separation package.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    # ==========================
+    # List models
+    # ==========================
+    list_parser = subparsers.add_parser(
+        "list",
+        help="List known models.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    list_parser.add_argument("--category", help="Filter by primary or secondary category.")
+    list_parser.add_argument("--all", action="store_true", help="Include models that are not supported for inference yet.")
+    list_parser.add_argument("--json", action="store_true")
+    list_parser.set_defaults(func=cmd_list)
+
+    # ==========================
+    # Show model info
+    # ==========================
+    info_parser = subparsers.add_parser(
+        "info",
+        help="Show model metadata.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    info_parser.add_argument("model")
+    info_parser.add_argument(
+        "--model-dir",
+        help="Local model cache directory. Defaults to PYMSS_MODEL_DIR, repository all_models if present, or ~/.cache/pymss/models.",
+    )
+    info_parser.set_defaults(func=cmd_info)
+
+    # ==========================
+    # Download models
+    # ==========================
+    download_parser = subparsers.add_parser(
+        "download",
+        help="Download a model by name, or use 'all'.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    download_parser.add_argument("model")
+    download_parser.add_argument(
+        "--model-dir",
+        help="Local model cache directory. Defaults to PYMSS_MODEL_DIR, repository all_models if present, or ~/.cache/pymss/models.",
+    )
+    download_parser.add_argument("--source", default="modelscope", choices=["modelscope", "huggingface", "hf-mirror"])
+    download_parser.add_argument("--endpoint", help="Custom resolve endpoint. It must serve files by relative path.")
+    download_parser.add_argument("--force", action="store_true")
+    download_parser.add_argument("--supported-only", action="store_true", help="Only used with model='all'.")
+    download_parser.set_defaults(func=cmd_download)
+
+    # ==========================
+    # Inference
+    # ==========================
+    infer_parser = subparsers.add_parser(
+        "infer",
+        help="Run inference by model name.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    infer_parser.add_argument("model")
+    infer_parser.add_argument(
+        "--model-dir",
+        help="Local model cache directory. Defaults to PYMSS_MODEL_DIR, repository all_models if present, or ~/.cache/pymss/models.",
+    )
+    infer_parser.add_argument(
+        "--download",
+        action="store_true",
+        help="Check/download the model before inference. Missing model files are downloaded automatically.",
+    )
+    infer_parser.add_argument("--source", default="modelscope", choices=["modelscope", "huggingface", "hf-mirror"])
+    infer_parser.add_argument("--endpoint", help="Custom resolve endpoint. It must serve files by relative path.")
+    infer_parser.add_argument("-i", "--input", required=True, help="Input audio file or folder.")
+    infer_parser.add_argument("-o", "--output", default="results", help="Output folder.")
+    infer_parser.add_argument(
+        "--save-as-folder",
+        action="store_true",
+        help="Save each input audio file's separated stems in a subfolder named after the audio file.",
+    )
+    infer_parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps", "mlx"])
+    infer_parser.add_argument(
+        "--device-id", action="append", type=int, dest="device_ids", help="CUDA device id. Can be repeated."
+    )
+    infer_parser.add_argument("--format", default="wav", choices=["wav", "flac", "mp3", "m4a"], dest="output_format")
+    infer_parser.add_argument("--wav-bit-depth", default="FLOAT", choices=["FLOAT", "PCM_16", "PCM_24"])
+    infer_parser.add_argument("--flac-bit-depth", default="PCM_16", choices=["PCM_16", "PCM_24"])
+    infer_parser.add_argument("--mp3-bit-rate", default="320k")
+    infer_parser.add_argument("--m4a-bit-rate", default="512k")
+    infer_parser.add_argument("--m4a-codec", default="aac")
+    infer_parser.add_argument("--m4a-aac-at-quality", default=2, type=int)
+    infer_parser.add_argument("--tta", action="store_true", help="Enable test time augmentation.")
+    infer_parser.add_argument("--debug", action="store_true")
+    infer_parser.add_argument(
+        "--param", action="append", default=[], help="Inference override as key=value, for example --param batch_size=2."
+    )
+    infer_parser.set_defaults(func=cmd_infer)
+
+    # ==========================
+    # Ensemble
+    # ==========================
+    ensemble_parser = subparsers.add_parser(
+        "ensemble",
+        help="Combine multiple audio files with an ensemble algorithm.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    ensemble_parser.add_argument("files", nargs="+", help="Input audio files. At least two files are required.")
+    ensemble_parser.add_argument(
+        "-a",
+        "--algorithm",
+        default="avg_wave",
+        choices=ENSEMBLE_ALGORITHMS,
+        help="Ensemble algorithm.",
+    )
+    ensemble_parser.add_argument(
+        "-w",
+        "--weights",
+        nargs="+",
+        type=float,
+        help="Input weights, for example --weights 1 0.8 1.2. Defaults to all 1.",
+    )
+    ensemble_parser.add_argument("-o", "--output", required=True, help="Output audio file.")
+    ensemble_parser.add_argument("--format", choices=["wav", "flac", "mp3", "m4a"], dest="output_format")
+    ensemble_parser.add_argument("--wav-bit-depth", default="FLOAT", choices=["FLOAT", "PCM_16", "PCM_24"])
+    ensemble_parser.add_argument("--flac-bit-depth", default="PCM_16", choices=["PCM_16", "PCM_24"])
+    ensemble_parser.add_argument("--mp3-bit-rate", default="320k")
+    ensemble_parser.add_argument("--m4a-bit-rate", default="512k")
+    ensemble_parser.add_argument("--m4a-codec", default="aac")
+    ensemble_parser.add_argument("--m4a-aac-at-quality", default=2, type=int)
+    ensemble_parser.set_defaults(func=cmd_ensemble)
+
+    # ==========================
+    # Workflow
+    # ==========================
+    workflow_parser = subparsers.add_parser(
+        "workflow",
+        help="Create, validate, or run an automatic multi-model workflow.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    workflow_subparsers = workflow_parser.add_subparsers(dest="workflow_command", required=True)
+
+    workflow_init_parser = workflow_subparsers.add_parser(
+        "init",
+        help="Write a starter workflow YAML file.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    workflow_init_parser.add_argument("-o", "--output", default="workflow.yaml", help="Workflow file to create.")
+    workflow_init_parser.add_argument("--force", action="store_true", help="Overwrite the output file if it exists.")
+    workflow_init_parser.set_defaults(func=cmd_workflow_init)
+
+    workflow_validate_parser = workflow_subparsers.add_parser(
+        "validate",
+        help="Validate a workflow YAML/JSON file.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    workflow_validate_parser.add_argument("-c", "--config", required=True, help="Workflow YAML/JSON file.")
+    workflow_validate_parser.add_argument(
+        "--model-dir",
+        help="Local model cache directory used when --require-files is set.",
+    )
+    workflow_validate_parser.add_argument(
+        "--check-models",
+        action="store_true",
+        help="Also check that every referenced model exists in the catalog.",
+    )
+    workflow_validate_parser.add_argument(
+        "--require-files",
+        action="store_true",
+        help="Also require every referenced catalog model file to exist locally.",
+    )
+    workflow_validate_parser.set_defaults(func=cmd_workflow_validate)
+
+    workflow_run_parser = workflow_subparsers.add_parser(
+        "run",
+        help="Run inference through a workflow YAML/JSON file.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    workflow_run_parser.add_argument("-c", "--config", required=True, help="Workflow YAML/JSON file.")
+    workflow_run_parser.add_argument("-i", "--input", required=True, help="Input audio file or folder.")
+    workflow_run_parser.add_argument("-o", "--output", default="results", help="Output folder.")
+    workflow_run_parser.add_argument(
+        "--output-layout",
+        default="folders",
+        choices=["folders", "flat"],
+        help=(
+            "Workflow output layout. 'folders' keeps each input under <output>/<audio>/; "
+            "'flat' writes outputs directly under the workflow task folder and save subfolders."
+        ),
+    )
+    workflow_run_parser.add_argument(
+        "--model-dir",
+        help="Local model cache directory. Workflow step model_dir values take precedence.",
+    )
+    workflow_run_parser.add_argument(
+        "--download",
+        action="store_true",
+        help="Download missing model files before each workflow step is loaded.",
+    )
+    workflow_run_parser.add_argument("--source", default="modelscope", choices=["modelscope", "huggingface", "hf-mirror"])
+    workflow_run_parser.add_argument("--endpoint", help="Custom resolve endpoint. It must serve files by relative path.")
+    workflow_run_parser.add_argument("--device", choices=["auto", "cpu", "cuda", "mps", "mlx"])
+    workflow_run_parser.add_argument("--format", choices=["wav", "flac", "mp3", "m4a"], dest="output_format")
+    workflow_run_parser.add_argument("--wav-bit-depth", default="FLOAT", choices=["FLOAT", "PCM_16", "PCM_24"])
+    workflow_run_parser.add_argument("--flac-bit-depth", default="PCM_16", choices=["PCM_16", "PCM_24"])
+    workflow_run_parser.add_argument("--mp3-bit-rate", default="320k")
+    workflow_run_parser.add_argument("--m4a-bit-rate", default="512k")
+    workflow_run_parser.add_argument("--m4a-codec", default="aac")
+    workflow_run_parser.add_argument("--m4a-aac-at-quality", default=2, type=int)
+    workflow_run_parser.add_argument("--debug", action="store_true")
+    workflow_run_parser.set_defaults(func=cmd_workflow_run)
+
+    # ==========================
+    # Server
+    # ==========================
+    serve_parser = subparsers.add_parser(
+        "serve",
+        help="Start an OpenAI-style HTTP inference server.",
+        formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60),
+    )
+    serve_parser.add_argument("model", nargs="?")
+    serve_parser.add_argument(
+        "--model-dir",
+        help="Local model cache directory. Defaults to PYMSS_MODEL_DIR, repository all_models if present, or ~/.cache/pymss/models.",
+    )
+    serve_parser.add_argument("--source", default="modelscope", choices=["modelscope", "huggingface", "hf-mirror"])
+    serve_parser.add_argument("--endpoint", help="Custom resolve endpoint. It must serve files by relative path.")
+    serve_parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps", "mlx"])
+    serve_parser.add_argument(
+        "--device-id",
+        action="append",
+        type=int,
+        dest="device_ids",
+        help="CUDA device id. Can be repeated.",
+    )
+    serve_parser.add_argument("--host", default="127.0.0.1")
+    serve_parser.add_argument("--port", default=8000, type=int)
+    serve_parser.add_argument("--api-key", help="Optional bearer token required for /v1/* endpoints.")
+    serve_parser.add_argument("--debug", action="store_true")
+    serve_parser.add_argument(
+        "--param",
+        action="append",
+        default=[],
+        help="Inference override as key=value, for example --param batch_size=2.",
+    )
+    serve_parser.add_argument("--max-audio-seconds", default=600.0, type=float)
+    serve_parser.add_argument("--max-request-bytes", default=536870912, type=int)
+    serve_parser.add_argument("--max-queue-size", default=8, type=int)
+    serve_parser.add_argument("--request-timeout-seconds", default=0.0, type=float)
+    serve_parser.add_argument("--webui", action="store_true", help="Serve the optional browser WebUI at /ui/.")
+    serve_parser.set_defaults(func=cmd_serve)
+
+    return parser
+
+
+def main(argv=None):
+    """Run the pymss command-line interface.
+
+    Args:
+        argv (Sequence[str] | None, optional): Command-line arguments. Uses sys.argv when None. Defaults to None.
+
+    Returns:
+        int: Process exit code."""
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        return args.func(args)
+    except Exception as exc:
+        print(f"pymss: error: {exc}", file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tools/pymss/config.py

@@ -0,0 +1,3 @@
+from pymss_core.config import AttrDict, ConfigLoader, load_config, to_attrdict, to_plain
+
+__all__ = ("AttrDict", "ConfigLoader", "load_config", "to_attrdict", "to_plain")

tools/pymss/ensemble.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import librosa
+import numpy as np
+
+from .audio_io import load_audio, save_audio
+
+ENSEMBLE_ALGORITHMS = (
+    "avg_wave",
+    "median_wave",
+    "min_wave",
+    "max_wave",
+    "avg_fft",
+    "median_fft",
+    "min_fft",
+    "max_fft",
+)
+
+
+def _as_channel_first(audio):
+    """Implement the as channel first helper.
+
+    Args:
+        audio (np.ndarray): Audio samples.
+
+    Returns:
+        Any: Computed result."""
+    audio = np.asarray(audio, dtype=np.float32)
+    return audio[None, :] if audio.ndim == 1 else audio
+
+
+def stft(wave, nfft=2048, hl=1024):
+    """Implement the stft helper.
+
+    Args:
+        wave (np.ndarray): Wave value.
+        nfft (Any, optional): Nfft value. Defaults to 2048.
+        hl (Any, optional): Hl value. Defaults to 1024.
+
+    Returns:
+        Any: Computed result."""
+    wave = _as_channel_first(wave)
+    return np.asfortranarray([librosa.stft(np.asfortranarray(channel), n_fft=nfft, hop_length=hl) for channel in wave])
+
+
+def istft(spec, hl=1024, length=None):
+    """Implement the istft helper.
+
+    Args:
+        spec (np.ndarray): Spec value.
+        hl (Any, optional): Hl value. Defaults to 1024.
+        length (Any, optional): Length value. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    return np.asfortranarray([librosa.istft(np.asfortranarray(channel), hop_length=hl, length=length) for channel in spec])
+
+
+def absmax(a, *, axis):
+    """Implement the absmax helper.
+
+    Args:
+        a (np.ndarray): A value.
+        axis (Any): Axis value.
+
+    Returns:
+        Any: Computed result."""
+    dims = list(a.shape)
+    dims.pop(axis)
+    indices = np.ogrid[tuple(slice(0, d) for d in dims)]
+    argmax = np.abs(a).argmax(axis=axis)
+    indices.insert((len(a.shape) + axis) % len(a.shape), argmax)
+    return a[tuple(indices)]
+
+
+def lambda_min(arr, axis=None, key=None, keepdims=False):
+    """Implement the lambda min helper.
+
+    Args:
+        arr (np.ndarray): Arr value.
+        axis (Any, optional): Axis value. Defaults to None.
+        key (str, optional): Key value. Defaults to None.
+        keepdims (Any, optional): Keepdims value. Defaults to False.
+
+    Returns:
+        Any: Computed result."""
+    idxs = np.argmin(key(arr), axis)
+    if axis is None:
+        return arr.flatten()[idxs]
+    idxs = np.expand_dims(idxs, axis)
+    result = np.take_along_axis(arr, idxs, axis)
+    return result if keepdims else np.squeeze(result, axis=axis)
+
+
+def lambda_max(arr, axis=None, key=None, keepdims=False):
+    """Implement the lambda max helper.
+
+    Args:
+        arr (np.ndarray): Arr value.
+        axis (Any, optional): Axis value. Defaults to None.
+        key (str, optional): Key value. Defaults to None.
+        keepdims (Any, optional): Keepdims value. Defaults to False.
+
+    Returns:
+        Any: Computed result."""
+    idxs = np.argmax(key(arr), axis)
+    if axis is None:
+        return arr.flatten()[idxs]
+    idxs = np.expand_dims(idxs, axis)
+    result = np.take_along_axis(arr, idxs, axis)
+    return result if keepdims else np.squeeze(result, axis=axis)
+
+
+def average_waveforms(pred_track, weights=None, algorithm="avg_wave"):
+    """Combine source waveforms with a selected ensemble algorithm.
+
+    Args:
+        pred_track (Any): Pred track value.
+        weights (Sequence[float] | None, optional): Per-file ensemble weights. Defaults to equal weights when None. Defaults to None.
+        algorithm (str, optional): Ensemble algorithm name. Defaults to "avg_wave".
+
+    Returns:
+        np.ndarray: Combined waveform shaped as channels by samples.
+
+    Example:
+        >>> combined = average_waveforms(predictions, weights=[1, 1], algorithm="avg_wave")"""
+    if algorithm not in ENSEMBLE_ALGORITHMS:
+        raise ValueError(f"Unknown ensemble algorithm: {algorithm}")
+
+    pred_track = np.asarray(pred_track, dtype=np.float32)
+    if pred_track.ndim != 3:
+        raise ValueError("pred_track must have shape (files, channels, samples)")
+
+    if weights is None:
+        weights = np.ones(pred_track.shape[0], dtype=np.float32)
+    weights = np.asarray(weights, dtype=np.float32)
+    if weights.shape != (pred_track.shape[0],):
+        raise ValueError("weights length must match number of input files")
+    if algorithm in {"avg_wave", "avg_fft"} and np.isclose(weights.sum(), 0.0):
+        raise ValueError("weights must not sum to zero for average ensemble algorithms")
+
+    final_length = pred_track.shape[-1]
+    mod_track = []
+    for idx in range(pred_track.shape[0]):
+        if algorithm == "avg_wave":
+            mod_track.append(pred_track[idx] * weights[idx])
+        elif algorithm in {"median_wave", "min_wave", "max_wave"}:
+            mod_track.append(pred_track[idx])
+        elif algorithm in {"avg_fft", "median_fft", "min_fft", "max_fft"}:
+            spec = stft(pred_track[idx], nfft=2048, hl=1024)
+            mod_track.append(spec * weights[idx] if algorithm == "avg_fft" else spec)
+
+    pred_track = np.asarray(mod_track)
+    if algorithm == "avg_wave":
+        return pred_track.sum(axis=0) / weights.sum()
+    if algorithm == "median_wave":
+        return np.median(pred_track, axis=0)
+    if algorithm == "min_wave":
+        return lambda_min(pred_track, axis=0, key=np.abs)
+    if algorithm == "max_wave":
+        return lambda_max(pred_track, axis=0, key=np.abs)
+    if algorithm == "avg_fft":
+        return istft(pred_track.sum(axis=0) / weights.sum(), hl=1024, length=final_length)
+    if algorithm == "min_fft":
+        return istft(lambda_min(pred_track, axis=0, key=np.abs), hl=1024, length=final_length)
+    if algorithm == "max_fft":
+        return istft(absmax(pred_track, axis=0), hl=1024, length=final_length)
+    if algorithm == "median_fft":
+        return istft(np.median(pred_track, axis=0), hl=1024, length=final_length)
+
+    raise AssertionError("unreachable")
+
+
+def ensemble_audios(files, algorithm="avg_wave", weights=None, logger=None):
+    """Load and combine multiple audio files with an ensemble algorithm.
+
+    All input files must have the same sample rate and channel count. If input
+    lengths differ, every file is truncated to the shortest length before
+    combining. ``avg_*`` algorithms use weights; median/min/max algorithms
+    ignore weights except for input validation.
+
+    Args:
+        files (Sequence[str | os.PathLike]): Audio files to combine. At least
+            two files are required.
+        algorithm (str, optional): Ensemble algorithm. Supported values are
+            ``avg_wave``, ``median_wave``, ``min_wave``, ``max_wave``,
+            ``avg_fft``, ``median_fft``, ``min_fft``, and ``max_fft``.
+            Defaults to ``"avg_wave"``.
+        weights (Sequence[float] | None, optional): Per-file weights. When
+            ``None``, every file gets weight ``1``. For average algorithms the
+            weight sum must not be zero. Defaults to None.
+        logger (logging.Logger | None, optional): Optional logger used for
+            debug messages and length-truncation warnings. Defaults to None.
+
+    Returns:
+        tuple[np.ndarray, int]: Combined audio shaped as samples by channels,
+        and the sample rate.
+
+    Raises:
+        ValueError: If fewer than two files are provided, weights length does
+            not match input count, sample rates differ, channel counts differ,
+            or the algorithm is unknown.
+        FileNotFoundError: If any input file does not exist.
+
+    Example:
+        >>> from pymss import ensemble_audios
+        >>> audio, sample_rate = ensemble_audios(
+        ...     ["vocals_a.wav", "vocals_b.wav"],
+        ...     algorithm="avg_wave",
+        ...     weights=[0.7, 0.3],
+        ... )
+
+    Example:
+        >>> audio, sample_rate = ensemble_audios(
+        ...     ["stem_a.wav", "stem_b.wav", "stem_c.wav"],
+        ...     algorithm="median_fft",
+        ... )"""
+    if len(files) < 2:
+        raise ValueError("at least two input files are required")
+
+    if weights is None:
+        weights = np.ones(len(files), dtype=np.float32)
+    weights = np.asarray(weights, dtype=np.float32)
+    if weights.shape != (len(files),):
+        raise ValueError("weights length must match number of input files")
+
+    data = []
+    sample_rate = None
+    for file in files:
+        path = Path(file)
+        if not path.is_file():
+            raise FileNotFoundError(f"input audio file not found: {path}")
+        audio, sr = load_audio(str(path), sr=None, mono=False)
+        audio = _as_channel_first(audio)
+        if sample_rate is None:
+            sample_rate = sr
+        elif sr != sample_rate:
+            raise ValueError(f"sample rate mismatch: {path} has {sr}, expected {sample_rate}")
+        data.append(audio)
+        if logger is not None:
+            logger.debug("read %s, waveform shape=%s, sample_rate=%s", path, audio.shape, sr)
+
+    channel_counts = {item.shape[0] for item in data}
+    if len(channel_counts) != 1:
+        raise ValueError("all input files must have the same channel count")
+
+    lengths = [item.shape[-1] for item in data]
+    min_length = min(lengths)
+    if len(set(lengths)) > 1:
+        if logger is not None:
+            logger.warning("Input audio files have different lengths. Truncating all to the shortest length.")
+        data = [item[..., :min_length] for item in data]
+
+    result = average_waveforms(np.asarray(data), weights=weights, algorithm=algorithm)
+    if logger is not None:
+        logger.debug("ensemble result shape=%s", result.shape)
+    return result.T, sample_rate
+
+
+def save_ensemble_audio(
+    files,
+    output,
+    algorithm="avg_wave",
+    weights=None,
+    output_format=None,
+    audio_params=None,
+    logger=None,
+):
+    """Combine audio files and save the ensemble result.
+
+    This is the file-writing wrapper around ``ensemble_audios(...)``. The
+    output format is inferred from ``output`` when it has a suffix, otherwise
+    ``output_format`` is used. If neither is provided, ``.wav`` is added.
+
+    Args:
+        files (Sequence[str | os.PathLike]): Audio files to combine. At least
+            two files are required.
+        output (str | os.PathLike): Output file path. A missing suffix becomes
+            ``.wav`` unless ``output_format`` is provided.
+        algorithm (str, optional): Ensemble algorithm. Supported values are
+            ``avg_wave``, ``median_wave``, ``min_wave``, ``max_wave``,
+            ``avg_fft``, ``median_fft``, ``min_fft``, and ``max_fft``.
+            Defaults to ``"avg_wave"``.
+        weights (Sequence[float] | None, optional): Per-file weights. Defaults
+            to equal weights when None.
+        output_format (str | None, optional): Explicit output format such as
+            ``wav``, ``flac``, ``mp3``, or ``m4a``. Defaults to None.
+        audio_params (dict | None, optional): Encoding options forwarded to
+            ``save_audio(...)``. Examples include
+            ``{"wav_bit_depth": "FLOAT"}``,
+            ``{"flac_bit_depth": "PCM_24"}``, or
+            ``{"mp3_bit_rate": "320k"}``. Defaults to None.
+        logger (logging.Logger | None, optional): Optional logger for progress
+            messages. Defaults to None.
+
+    Returns:
+        pathlib.Path: Final output path.
+
+    Raises:
+        ValueError: If inputs cannot be ensembled.
+        FileNotFoundError: If any input file does not exist.
+
+    Example:
+        >>> from pymss import save_ensemble_audio
+        >>> save_ensemble_audio(
+        ...     ["vocals_a.wav", "vocals_b.wav"],
+        ...     "vocals_ensemble.flac",
+        ...     algorithm="avg_wave",
+        ...     weights=[1, 1],
+        ...     audio_params={"flac_bit_depth": "PCM_24"},
+        ... )
+
+    Example:
+        >>> save_ensemble_audio(["a.wav", "b.wav"], "ensemble", output_format="wav")"""
+    result, sample_rate = ensemble_audios(files, algorithm=algorithm, weights=weights, logger=logger)
+    output_path = Path(output)
+    if not output_path.suffix and not output_format:
+        output_path = output_path.with_suffix(".wav")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_format = output_format or output_path.suffix.lstrip(".").lower() or "wav"
+    save_audio(str(output_path), result, sample_rate, output_format, audio_params or {})
+    return output_path

tools/pymss/logger.py

@@ -0,0 +1,339 @@
+import gzip
+import logging
+import os
+import shutil
+import sys
+from datetime import datetime
+
+MAX_LOG = 100
+LOG_DIR = ".logs"
+LOG_ENV_NAME = "PYMSS_LOG_FILE"
+
+
+def _safe_relpath(pathname):
+    """Implement the safe relpath helper.
+
+    Args:
+        pathname (str): Pathname value.
+
+    Returns:
+        Any: Computed result."""
+    if not pathname:
+        return pathname
+
+    normalized = os.path.normpath(pathname)
+    if os.name == "nt" and normalized.startswith("\\\\?\\"):
+        normalized = normalized[4:]
+
+    try:
+        return os.path.relpath(normalized)
+    except ValueError:
+        return normalized
+
+
+class ColorFormatter(logging.Formatter):
+    """Console log formatter with optional ANSI colors.
+
+    Args:
+        enable_color (Any, optional): Enable color value. Defaults to True.
+    """
+
+    COLORS = {
+        "DBG": "\033[1;36m",
+        "INF": "\033[1;32m",
+        "WAR": "\033[1;33m",
+        "ERR": "\033[1;31m",
+        "CRI": "\033[1;35m",
+    }
+    MESSAGE_COLORS = {
+        "DBG": "\033[36m",
+        "INF": "\033[32m",
+        "WAR": "\033[33m",
+        "ERR": "\033[31m",
+        "CRI": "\033[35m",
+    }
+    RESET = "\033[0m"
+    LEVEL_MAP = {
+        "DEBUG": "DBG",
+        "INFO": "INF",
+        "WARNING": "WAR",
+        "ERROR": "ERR",
+        "CRITICAL": "CRI",
+    }
+
+    def __init__(self, enable_color=True):
+        """Initialize the instance.
+
+        Args:
+            enable_color (Any, optional): Enable color value. Defaults to True.
+
+        Returns:
+            None: This method completes for its side effects."""
+        super().__init__(
+            fmt="%(asctime)s | %(levelname)s | %(pathname)s:%(lineno)d | %(message)s",
+            datefmt="%H:%M:%S",
+        )
+        self.enable_color = enable_color
+
+    def format(self, record):
+        """Format value.
+
+        Args:
+            record (Any): Record value.
+
+        Returns:
+            Any: Computed result."""
+        record.pathname = _safe_relpath(record.pathname)
+        original_levelname = record.levelname
+        original_msg = record.msg
+
+        short_level = self.LEVEL_MAP.get(record.levelname, record.levelname[:3])
+        if self.enable_color:
+            level_color = self.COLORS.get(short_level, "")
+            message_color = self.MESSAGE_COLORS.get(short_level, "")
+            record.levelname = f"{level_color}{short_level}{self.RESET}"
+            record.msg = f"{message_color}{record.getMessage()}{self.RESET}"
+            record.args = ()
+        else:
+            record.levelname = short_level
+
+        try:
+            return super().format(record)
+        finally:
+            record.levelname = original_levelname
+            record.msg = original_msg
+
+
+class FileFormatter(logging.Formatter):
+    """File log formatter that writes stable relative paths."""
+
+    LEVEL_MAP = ColorFormatter.LEVEL_MAP
+
+    def __init__(self):
+        """Initialize the instance.
+
+        Args:
+            None: This callable does not accept user-provided arguments.
+
+        Returns:
+            None: This method completes for its side effects."""
+        super().__init__(
+            fmt="%(asctime)s | %(levelname)s | %(pathname)s:%(lineno)d | %(message)s",
+            datefmt="%H:%M:%S",
+        )
+
+    def format(self, record):
+        """Format value.
+
+        Args:
+            record (Any): Record value.
+
+        Returns:
+            Any: Computed result."""
+        record.pathname = _safe_relpath(record.pathname)
+        original_levelname = record.levelname
+        record.levelname = self.LEVEL_MAP.get(record.levelname, record.levelname[:3])
+        try:
+            return super().format(record)
+        finally:
+            record.levelname = original_levelname
+
+
+def _supports_color(stream):
+    """Implement the supports color helper.
+
+    Args:
+        stream (Any): Stream value.
+
+    Returns:
+        Any: Computed result."""
+    if os.environ.get("NO_COLOR"):
+        return False
+    if os.environ.get("PYMSS_FORCE_COLOR"):
+        return True
+    return hasattr(stream, "isatty") and stream.isatty()
+
+
+def _compress_log_file(path):
+    """Implement the compress log file helper.
+
+    Args:
+        path (str | os.PathLike): File system path.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    gz_path = f"{path}.gz"
+    if os.path.exists(gz_path):
+        return
+    try:
+        with open(path, "rb") as src, gzip.open(gz_path, "wb") as dst:
+            shutil.copyfileobj(src, dst)
+        os.remove(path)
+    except OSError:
+        pass
+
+
+def _parse_log_time(filename):
+    """Parse log time.
+
+    Args:
+        filename (str): Filename value.
+
+    Returns:
+        Any: Parsed value."""
+    stem = filename
+    if stem.endswith(".gz"):
+        stem = stem[:-3]
+    if stem.endswith(".log"):
+        stem = stem[:-4]
+    for fmt in ("%Y-%m-%d_%H-%M-%S", "%Y-%m-%d"):
+        try:
+            return datetime.strptime(stem, fmt)
+        except ValueError:
+            continue
+    return datetime.min
+
+
+def manage_log_files(log_dir, max_log):
+    """Compress or remove old log files according to the retention limit.
+
+    Args:
+        log_dir (Any): Log dir value.
+        max_log (int): Max log value.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    try:
+        log_files = [filename for filename in os.listdir(log_dir) if filename.endswith(".log") or filename.endswith(".log.gz")]
+    except OSError:
+        return
+
+    current_log = os.environ.get(LOG_ENV_NAME)
+    for filename in log_files:
+        path = os.path.join(log_dir, filename)
+        if filename.endswith(".log") and path != current_log:
+            _compress_log_file(path)
+
+    try:
+        log_files = [filename for filename in os.listdir(log_dir) if filename.endswith(".log") or filename.endswith(".log.gz")]
+    except OSError:
+        return
+
+    log_files = sorted(log_files, key=_parse_log_time)
+    while len(log_files) > max_log:
+        oldest_file = log_files.pop(0)
+        path = os.path.join(log_dir, oldest_file)
+        if path == current_log:
+            continue
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+
+def _get_log_path(log_dir):
+    """Return log path.
+
+    Args:
+        log_dir (Any): Log dir value.
+
+    Returns:
+        Any: Computed result."""
+    log_path = os.environ.get(LOG_ENV_NAME)
+    if log_path:
+        return log_path
+
+    os.makedirs(log_dir, exist_ok=True)
+    log_path = os.path.join(log_dir, datetime.now().strftime("%Y-%m-%d_%H-%M-%S.log"))
+    os.environ[LOG_ENV_NAME] = log_path
+    return log_path
+
+
+def set_log_level(logger, level):
+    """Set the level for every handler attached to a logger.
+
+    Args:
+        logger (logging.Logger | None): Optional logger for progress messages.
+        level (int | str): Level value.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    if hasattr(logger, "console_handler"):
+        logger.console_handler.setLevel(level)
+
+
+def get_separation_logger(
+    console_level=logging.INFO,
+    enable_file_log=False,
+    max_log=MAX_LOG,
+    log_dir=LOG_DIR,
+    enable_color=None,
+):
+    """Create or return the shared pymss separation logger.
+
+    The logger writes concise console messages by default and can optionally
+    add a debug-level file handler. Calling this function repeatedly returns
+    the same logger and updates the console handler level.
+
+    Args:
+        console_level (int, optional): Console handler log level, such as
+            ``logging.INFO`` or ``logging.DEBUG``. Defaults to logging.INFO.
+        enable_file_log (bool, optional): Whether to add a file handler. File
+            logs are written at debug level regardless of ``console_level``.
+            Defaults to False.
+        max_log (int, optional): Maximum number of ``.log`` or ``.log.gz``
+            files to keep in ``log_dir``. Older files are compressed or
+            removed. Defaults to ``MAX_LOG``.
+        log_dir (str | os.PathLike, optional): Directory for file logs when
+            ``enable_file_log`` is true. Defaults to ``LOG_DIR``.
+        enable_color (bool | None, optional): Whether console output should use
+            ANSI colors. ``None`` auto-detects color support and respects
+            ``NO_COLOR``/``PYMSS_FORCE_COLOR``. Defaults to None.
+
+    Returns:
+        logging.Logger: Shared pymss logger. The object also stores
+        ``console_handler`` and, when enabled, ``file_handler`` attributes.
+
+    Example:
+        >>> import logging
+        >>> from pymss import get_separation_logger
+        >>> logger = get_separation_logger(console_level=logging.DEBUG)
+        >>> logger.debug("debug output is visible")
+
+    Example:
+        >>> logger = get_separation_logger(enable_file_log=True, log_dir=".logs")
+        >>> logger.info("message is written to console and log file")"""
+    logger = logging.getLogger("logger")
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+
+    if hasattr(logger, "console_handler"):
+        logger.console_handler.setLevel(console_level)
+        if enable_file_log and not hasattr(logger, "file_handler"):
+            log_path = _get_log_path(log_dir)
+            file_handler = logging.FileHandler(log_path, mode="a", encoding="utf-8")
+            file_handler.setLevel(logging.DEBUG)
+            file_handler.setFormatter(FileFormatter())
+            logger.addHandler(file_handler)
+            logger.file_handler = file_handler
+            manage_log_files(log_dir, max_log)
+        return logger
+
+    console_handler = logging.StreamHandler(sys.stderr)
+    console_handler.setLevel(console_level)
+    if enable_color is None:
+        enable_color = _supports_color(console_handler.stream)
+    console_handler.setFormatter(ColorFormatter(enable_color=enable_color))
+    logger.addHandler(console_handler)
+    logger.console_handler = console_handler
+
+    if enable_file_log:
+        log_path = _get_log_path(log_dir)
+        file_handler = logging.FileHandler(log_path, mode="a", encoding="utf-8")
+        file_handler.setLevel(logging.DEBUG)
+        file_handler.setFormatter(FileFormatter())
+        logger.addHandler(file_handler)
+        logger.file_handler = file_handler
+        manage_log_files(log_dir, max_log)
+
+    return logger

tools/pymss/model_download.py

@@ -0,0 +1,399 @@
+import hashlib
+import json
+import os
+import shutil
+import subprocess
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from pathlib import Path
+
+from tqdm import tqdm
+
+from .model_registry import (
+    auxiliary_paths_for,
+    config_path_for,
+    get_model_entry,
+    model_path_for,
+)
+
+
+HF_REPO = "baicai1145/pymss"
+MS_REPO = "baicai1145/pymss"
+HF_BASE_URL = f"https://huggingface.co/{HF_REPO}/resolve/main"
+MS_BASE_URL = f"https://www.modelscope.cn/models/{MS_REPO}/resolve/master"
+MS_FILES_API = f"https://www.modelscope.cn/api/v1/models/{MS_REPO}/repo/files?Revision=master&Recursive=true"
+MODEL_FILE_SUFFIXES = {".ckpt", ".th", ".pth", ".chpt", ".safetensors", ".pt", ".yaml", ".yml", ".json"}
+ARIA2C_PATH = shutil.which("aria2c")
+
+
+class DownloadError(RuntimeError):
+    """Base exception raised when model download fails."""
+
+    pass
+
+
+class DownloadValidationError(DownloadError):
+    """Exception raised when a downloaded file fails size or hash validation."""
+
+    pass
+
+
+def _quote_path(path):
+    """Implement the quote path helper.
+
+    Args:
+        path (str | os.PathLike): File system path.
+
+    Returns:
+        Any: Computed result."""
+    return urllib.parse.quote(path, safe="/")
+
+
+def remote_url(relpath, source="modelscope", endpoint=None):
+    """Build the remote download URL for a catalog-relative file path.
+
+    Args:
+        relpath (str): Relpath value.
+        source (str, optional): Download source name. Defaults to "modelscope".
+        endpoint (str | None, optional): Optional custom download endpoint. Defaults to None.
+
+    Returns:
+        str: Absolute URL for the requested source."""
+    if endpoint:
+        return f"{endpoint.rstrip('/')}/{_quote_path(relpath)}"
+    if source == "huggingface":
+        return f"{HF_BASE_URL}/{_quote_path(relpath)}"
+    if source == "hf-mirror":
+        return f"https://hf-mirror.com/{HF_REPO}/resolve/main/{_quote_path(relpath)}"
+    if source == "modelscope":
+        return f"{MS_BASE_URL}/{_quote_path(relpath)}"
+    raise ValueError("source must be one of: modelscope, huggingface, hf-mirror")
+
+
+def _read_json_url(url, timeout=30):
+    """Read json url.
+
+    Args:
+        url (str): Url value.
+        timeout (int, optional): Network timeout in seconds. Defaults to 30.
+
+    Returns:
+        Any: Computed result."""
+    with urllib.request.urlopen(url, timeout=timeout) as response:
+        return json.load(response)
+
+
+def fetch_modelscope_file_index(timeout=30):
+    """Fetch modelscope file index.
+
+    Args:
+        timeout (int, optional): Network timeout in seconds. Defaults to 30.
+
+    Returns:
+        Any: Computed result."""
+    data = _read_json_url(MS_FILES_API, timeout=timeout)
+    files = data.get("Data", {}).get("Files", [])
+    return {item["Path"]: item for item in files if item.get("Type") == "blob"}
+
+
+def _sha256(path):
+    """Implement the sha256 helper.
+
+    Args:
+        path (str | os.PathLike): File system path.
+
+    Returns:
+        Any: Computed result."""
+    digest = hashlib.sha256()
+    with open(path, "rb") as f:
+        while True:
+            chunk = f.read(1024 * 1024)
+            if not chunk:
+                break
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _expected_size_and_hash(relpath, source_index):
+    """Implement the expected size and hash helper.
+
+    Args:
+        relpath (str): Relpath value.
+        source_index (Any): Source index value.
+
+    Returns:
+        Any: Computed result."""
+    if not source_index:
+        return None, ""
+    item = source_index.get(relpath, {})
+    size = item.get("Size")
+    sha256 = item.get("Sha256") or ""
+    return int(size) if size else None, sha256
+
+
+def _already_valid(path, expected_size=None, expected_sha256=""):
+    """Implement the already valid helper.
+
+    Args:
+        path (str | os.PathLike): File system path.
+        expected_size (Any, optional): Expected size value. Defaults to None.
+        expected_sha256 (Any, optional): Expected sha256 value. Defaults to ''.
+
+    Returns:
+        Any: Computed result."""
+    if not path.is_file():
+        return False
+    if expected_size is not None and path.stat().st_size != expected_size:
+        return False
+    if expected_sha256 and _sha256(path) != expected_sha256:
+        return False
+    return True
+
+
+def _cleanup_partial_download(tmp):
+    """Implement the cleanup partial download helper.
+
+    Args:
+        tmp (Any): Tmp value.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    for path in (tmp, Path(str(tmp) + ".aria2")):
+        try:
+            path.unlink()
+        except FileNotFoundError:
+            pass
+
+
+def _validate_downloaded_file(path, dest, expected_size=None, expected_sha256=""):
+    """Validate downloaded file.
+
+    Args:
+        path (str | os.PathLike): File system path.
+        dest (Any): Dest value.
+        expected_size (Any, optional): Expected size value. Defaults to None.
+        expected_sha256 (Any, optional): Expected sha256 value. Defaults to ''.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    if expected_size is not None and path.stat().st_size != expected_size:
+        raise DownloadValidationError(f"size mismatch for {dest.name}: expected {expected_size}, got {path.stat().st_size}")
+    if expected_sha256:
+        actual = _sha256(path)
+        if actual != expected_sha256:
+            raise DownloadValidationError(f"sha256 mismatch for {dest.name}: expected {expected_sha256}, got {actual}")
+
+
+def _download_file_urllib(url, tmp, dest, expected_size=None, expected_sha256="", timeout=30):
+    """Download file urllib.
+
+    Args:
+        url (str): Url value.
+        tmp (Any): Tmp value.
+        dest (Any): Dest value.
+        expected_size (Any, optional): Expected size value. Defaults to None.
+        expected_sha256 (Any, optional): Expected sha256 value. Defaults to ''.
+        timeout (int, optional): Network timeout in seconds. Defaults to 30.
+
+    Returns:
+        Any: Computed result."""
+    with urllib.request.urlopen(url, timeout=timeout) as response:
+        total = int(response.headers.get("content-length") or expected_size or 0)
+        with open(tmp, "wb") as f, tqdm(total=total, unit="B", unit_scale=True, desc=dest.name) as progress:
+            while True:
+                chunk = response.read(1024 * 1024)
+                if not chunk:
+                    break
+                f.write(chunk)
+                progress.update(len(chunk))
+    _validate_downloaded_file(tmp, dest, expected_size, expected_sha256)
+    os.replace(tmp, dest)
+    return dest
+
+
+def _download_file_aria2(url, tmp, dest, expected_size=None, expected_sha256="", timeout=30):
+    """Download file aria2.
+
+    Args:
+        url (str): Url value.
+        tmp (Any): Tmp value.
+        dest (Any): Dest value.
+        expected_size (Any, optional): Expected size value. Defaults to None.
+        expected_sha256 (Any, optional): Expected sha256 value. Defaults to ''.
+        timeout (int, optional): Network timeout in seconds. Defaults to 30.
+
+    Returns:
+        Any: Computed result."""
+    cmd = [
+        ARIA2C_PATH,
+        "--allow-overwrite=true",
+        "--auto-file-renaming=false",
+        "--continue=true",
+        "--console-log-level=warn",
+        "--summary-interval=1",
+        "--max-connection-per-server=16",
+        "--split=16",
+        "--min-split-size=1M",
+        "--max-tries=3",
+        f"--connect-timeout={timeout}",
+        f"--timeout={timeout}",
+        "--dir",
+        str(tmp.parent),
+        "--out",
+        tmp.name,
+        url,
+    ]
+    result = subprocess.run(cmd, check=False)
+    if result.returncode != 0:
+        raise DownloadError(f"aria2c failed with exit code {result.returncode}")
+    if not tmp.is_file():
+        raise DownloadError("aria2c did not create the expected output file")
+
+    _validate_downloaded_file(tmp, dest, expected_size, expected_sha256)
+    os.replace(tmp, dest)
+    return dest
+
+
+def _download_file(url, dest, expected_size=None, expected_sha256="", timeout=30, retries=2):
+    """Download file.
+
+    Args:
+        url (str): Url value.
+        dest (Any): Dest value.
+        expected_size (Any, optional): Expected size value. Defaults to None.
+        expected_sha256 (Any, optional): Expected sha256 value. Defaults to ''.
+        timeout (int, optional): Network timeout in seconds. Defaults to 30.
+        retries (int, optional): Retries value. Defaults to 2.
+
+    Returns:
+        Any: Computed result."""
+    dest = Path(dest)
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_name(dest.name + ".part")
+    last_error = None
+
+    for attempt in range(retries + 1):
+        try:
+            if ARIA2C_PATH:
+                return _download_file_aria2(url, tmp, dest, expected_size, expected_sha256, timeout=timeout)
+            return _download_file_urllib(url, tmp, dest, expected_size, expected_sha256, timeout=timeout)
+        except (OSError, urllib.error.URLError, urllib.error.HTTPError, DownloadError) as exc:
+            last_error = exc
+            # don't rm temp file if aria2c is being used
+            if not ARIA2C_PATH or isinstance(exc, DownloadValidationError):
+                _cleanup_partial_download(tmp)
+            if attempt < retries:
+                time.sleep(1.0 + attempt)
+    raise DownloadError(f"failed to download {url}: {last_error}")
+
+
+def files_for_model(model_name, model_dir=None):
+    """Return the local file targets required by a catalog model.
+
+    Args:
+        model_name (str): Model name or alias from the pymss catalog.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        tuple[ModelEntry, list[tuple[str, Path]]]: Catalog entry and required files."""
+    entry = get_model_entry(model_name)
+    files = [(entry.relpath, model_path_for(entry, model_dir))]
+    config_path = config_path_for(entry, model_dir)
+    if entry.config_relpath and config_path is not None:
+        files.append((entry.config_relpath, config_path))
+    files.extend(zip(entry.auxiliary_relpaths, auxiliary_paths_for(entry, model_dir)))
+    return entry, files
+
+
+def download_model(model_name, model_dir=None, source="modelscope", endpoint=None, verify=True, force=False, timeout=30):
+    """Download all files required by one catalog model.
+
+    The downloader resolves the model from the pymss catalog, downloads the
+    weights, config, and auxiliary files, and skips files that already match
+    available size/hash metadata. If ``aria2c`` is available on ``PATH``, pymss
+    uses it for resumable multi-connection downloads; otherwise it falls back
+    to Python's urllib downloader.
+
+    Args:
+        model_name (str): Model name, stem, or alias from the pymss catalog.
+        model_dir (str | os.PathLike | None, optional): Local model cache
+            directory. When omitted, pymss uses its default model directory.
+            Defaults to None.
+        source (str, optional): Download source. Supported values are
+            ``modelscope``, ``huggingface``, and ``hf-mirror``. Defaults to
+            ``"modelscope"``.
+        endpoint (str | None, optional): Custom endpoint prefix. When provided,
+            the final URL is ``endpoint/relative/catalog/path`` and ``source``
+            is ignored for URL construction. Defaults to None.
+        verify (bool, optional): Whether to validate downloads with available
+            ModelScope size/hash metadata. Custom endpoints skip source index
+            lookup. Defaults to True.
+        force (bool, optional): Whether to redownload files even when existing
+            local files appear valid. Defaults to False.
+        timeout (int, optional): Network timeout in seconds. Defaults to 30.
+
+    Returns:
+        dict: Download result with ``entry`` (``ModelEntry``), ``downloaded``
+        (list of paths written this call), and ``skipped`` (list of existing
+        valid paths).
+
+    Raises:
+        KeyError: If ``model_name`` is unknown.
+        DownloadError: If downloading fails after retries.
+        DownloadValidationError: If a downloaded file fails size/hash checks.
+
+    Example:
+        >>> from pymss import download_model
+        >>> result = download_model("bs_roformer_voc_hyperacev2", model_dir="models")
+        >>> result["entry"].name
+
+    Example:
+        >>> download_model(
+        ...     "bs_roformer_voc_hyperacev2",
+        ...     source="hf-mirror",
+        ...     force=True,
+        ...     timeout=60,
+        ... )"""
+    entry, files = files_for_model(model_name, model_dir)
+    index = fetch_modelscope_file_index(timeout=timeout) if verify and endpoint is None else None
+    downloaded = []
+    skipped = []
+
+    for relpath, dest in files:
+        expected_size, expected_sha256 = _expected_size_and_hash(relpath, index)
+        if not force and _already_valid(dest, expected_size, expected_sha256):
+            skipped.append(str(dest))
+            continue
+        url = remote_url(relpath, source=source, endpoint=endpoint)
+        _download_file(url, dest, expected_size, expected_sha256, timeout=timeout)
+        downloaded.append(str(dest))
+
+    return {"entry": entry, "downloaded": downloaded, "skipped": skipped}
+
+
+def download_all(model_dir=None, source="modelscope", endpoint=None, supported_only=False, force=False, timeout=30):
+    """Download every catalog model, optionally limited to supported entries.
+
+    Args:
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+        source (str, optional): Download source name. Defaults to "modelscope".
+        endpoint (str | None, optional): Optional custom download endpoint. Defaults to None.
+        supported_only (bool, optional): Supported only value. Defaults to False.
+        force (bool, optional): Whether to overwrite or redownload existing files. Defaults to False.
+        timeout (int, optional): Network timeout in seconds. Defaults to 30.
+
+    Returns:
+        list[dict]: Per-model download results."""
+    from .model_registry import list_models
+
+    results = []
+    for entry in list_models(supported=True if supported_only else None):
+        try:
+            results.append(
+                download_model(entry.name, model_dir=model_dir, source=source, endpoint=endpoint, force=force, timeout=timeout)
+            )
+        except Exception as exc:
+            results.append({"entry": entry, "error": str(exc)})
+    return results

tools/pymss/model_registry.py

@@ -0,0 +1,386 @@
+import json
+import os
+from dataclasses import dataclass
+from functools import lru_cache
+from importlib import resources
+from pathlib import Path
+
+
+def _default_model_dir():
+    """Implement the default model dir helper.
+
+    Args:
+        None: This callable does not accept user-provided arguments.
+
+    Returns:
+        Any: Computed result."""
+    env_value = os.environ.get("PYMSS_MODEL_DIR")
+    if env_value:
+        return Path(env_value)
+    repo_models = Path(__file__).resolve().parent.parent / "all_models"
+    if repo_models.is_dir():
+        return repo_models
+    return Path.home() / ".cache" / "pymss" / "models"
+
+
+DEFAULT_MODEL_DIR = _default_model_dir()
+
+
+@dataclass(frozen=True)
+class ModelEntry:
+    """Catalog metadata for one downloadable pymss model."""
+
+    name: str
+    aliases: tuple
+    model_type: str | None
+    architecture: str
+    supported: bool
+    unsupported_reason: str
+    relpath: str
+    config_relpath: str
+    auxiliary_relpaths: tuple
+    size_bytes: int
+    sha256: str
+    primary_category: str
+    primary_category_cn: str
+    secondary_category: str
+    secondary_category_cn: str
+    target_stem: str
+    config_instruments: str
+    config_target_instrument: str
+    classification_confidence: str
+    classification_basis: str
+
+    @property
+    def stem(self):
+        """Implement the stem helper.
+
+        Args:
+            None: This callable does not accept user-provided arguments.
+
+        Returns:
+            Any: Computed result."""
+        return Path(self.name).stem
+
+    @property
+    def category_path(self):
+        """Implement the category path helper.
+
+        Args:
+            None: This callable does not accept user-provided arguments.
+
+        Returns:
+            Any: Computed result."""
+        return "/".join(part for part in (self.primary_category, self.secondary_category) if part)
+
+    @classmethod
+    def from_dict(cls, data):
+        """Implement the from dict helper.
+
+        Args:
+            data (Mapping | None): Data value.
+
+        Returns:
+            Any: Computed result."""
+        return cls(
+            name=data["name"],
+            aliases=tuple(data.get("aliases", ())),
+            model_type=data.get("model_type"),
+            architecture=data.get("architecture", ""),
+            supported=bool(data.get("supported", False)),
+            unsupported_reason=data.get("unsupported_reason", ""),
+            relpath=data["relpath"],
+            config_relpath=data.get("config_relpath", ""),
+            auxiliary_relpaths=tuple(data.get("auxiliary_relpaths", ())),
+            size_bytes=int(data.get("size_bytes", 0)),
+            sha256=data.get("sha256", ""),
+            primary_category=data.get("primary_category", ""),
+            primary_category_cn=data.get("primary_category_cn", ""),
+            secondary_category=data.get("secondary_category", ""),
+            secondary_category_cn=data.get("secondary_category_cn", ""),
+            target_stem=data.get("target_stem", ""),
+            config_instruments=data.get("config_instruments", ""),
+            config_target_instrument=data.get("config_target_instrument", ""),
+            classification_confidence=data.get("classification_confidence", ""),
+            classification_basis=data.get("classification_basis", ""),
+        )
+
+
+@lru_cache(maxsize=1)
+def load_model_catalog():
+    """Load model catalog.
+
+    Args:
+        None: This callable does not accept user-provided arguments.
+
+    Returns:
+        Any: Computed result."""
+    with resources.files("pymss.resources").joinpath("model_catalog.json").open(encoding="utf-8") as f:
+        data = json.load(f)
+    models = [ModelEntry.from_dict(item) for item in data["models"]]
+    return {**data, "models": models}
+
+
+@lru_cache(maxsize=1)
+def _model_index():
+    """Implement the model index helper.
+
+    Args:
+        None: This callable does not accept user-provided arguments.
+
+    Returns:
+        Any: Computed result."""
+    index = {}
+    for entry in load_model_catalog()["models"]:
+        names = {entry.name, entry.stem, *entry.aliases}
+        for name in names:
+            key = _normalize_model_name(name)
+            if key in index and index[key].name != entry.name:
+                continue
+            index[key] = entry
+    return index
+
+
+def _normalize_model_name(name):
+    """Normalize model name.
+
+    Args:
+        name (Any): Name value.
+
+    Returns:
+        Any: Computed result."""
+    return str(name).strip().lower()
+
+
+def list_models(category=None, supported=None):
+    """List model catalog entries.
+
+    The catalog contains every model known to pymss, including unsupported
+    entries. Use the filters when building model selectors, download tools, or
+    validation code.
+
+    Args:
+        category (str | None, optional): Optional category filter. The value is
+            matched against primary category, secondary category, or combined
+            ``primary/secondary`` category path. Matching is case-insensitive.
+            Defaults to None.
+        supported (bool | None, optional): Support-status filter. ``True``
+            returns only models supported by the current inference code,
+            ``False`` returns unsupported entries, and ``None`` returns all
+            catalog entries. Defaults to None.
+
+    Returns:
+        list[ModelEntry]: Matching catalog entries in catalog order.
+
+    Example:
+        >>> from pymss import list_models
+        >>> supported_models = list_models(supported=True)
+        >>> supported_models[0].name
+
+    Example:
+        >>> vocal_models = list_models(category="vocal", supported=True)
+        >>> [model.stem for model in vocal_models[:3]]"""
+    models = load_model_catalog()["models"]
+    if category:
+        category = category.lower()
+        models = [
+            item
+            for item in models
+            if item.primary_category.lower() == category
+            or item.secondary_category.lower() == category
+            or item.category_path.lower() == category
+        ]
+    if supported is not None:
+        models = [item for item in models if item.supported is bool(supported)]
+    return models
+
+
+def get_model_entry(model_name):
+    """Return catalog metadata for one model name or alias.
+
+    Args:
+        model_name (str): Full catalog filename, stem name, or alias. Matching
+            is case-insensitive after stripping surrounding whitespace.
+
+    Returns:
+        ModelEntry: Catalog entry containing architecture, support status,
+        relative file paths, hashes, categories, target stem, and aliases.
+
+    Raises:
+        KeyError: If ``model_name`` is unknown.
+
+    Example:
+        >>> from pymss import get_model_entry
+        >>> entry = get_model_entry("bs_roformer_voc_hyperacev2")
+        >>> entry.model_type
+        'bs_roformer'
+
+    Example:
+        >>> entry.supported, entry.category_path
+        (True, entry.category_path)"""
+    try:
+        return _model_index()[_normalize_model_name(model_name)]
+    except KeyError as exc:
+        raise KeyError(f"Unknown pymss model: {model_name}") from exc
+
+
+def model_root(model_dir=None):
+    """Implement the model root helper.
+
+    Args:
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    return Path(model_dir).expanduser() if model_dir else DEFAULT_MODEL_DIR
+
+
+def model_path_for(entry, model_dir=None):
+    """Implement the model path for helper.
+
+    Args:
+        entry (ModelEntry): Entry value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    return model_root(model_dir) / entry.relpath
+
+
+def config_path_for(entry, model_dir=None):
+    """Implement the config path for helper.
+
+    Args:
+        entry (ModelEntry): Entry value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    return model_root(model_dir) / entry.config_relpath if entry.config_relpath else None
+
+
+def auxiliary_paths_for(entry, model_dir=None):
+    """Implement the auxiliary paths for helper.
+
+    Args:
+        entry (ModelEntry): Entry value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    root = model_root(model_dir)
+    return [root / relpath for relpath in entry.auxiliary_relpaths]
+
+
+def resolve_model(model_name, model_dir=None, require_supported=True, require_exists=True):
+    """Resolve a catalog model to local file paths.
+
+    This function does not instantiate a model. It only translates a catalog
+    name or alias into the local weights/config paths that ``MSSeparator`` will
+    use.
+
+    Args:
+        model_name (str): Model name, stem, or alias from the pymss catalog.
+        model_dir (str | os.PathLike | None, optional): Local model cache
+            directory. When omitted, pymss uses ``PYMSS_MODEL_DIR`` if set, a
+            repository-local ``all_models`` directory if present, or the user
+            cache under ``~/.cache/pymss/models``. Defaults to None.
+        require_supported (bool, optional): Whether unsupported catalog entries
+            should raise ``ValueError``. Defaults to True.
+        require_exists (bool, optional): Whether resolved model, config, and
+            auxiliary files must already exist locally. Defaults to True.
+
+    Returns:
+        dict: Dictionary with ``entry`` (``ModelEntry``), ``model_type``,
+        ``model_path``, and ``config_path`` keys.
+
+    Raises:
+        KeyError: If the model name is unknown.
+        ValueError: If the model is unsupported and ``require_supported`` is
+            true.
+        FileNotFoundError: If required local files are missing and
+            ``require_exists`` is true.
+
+    Example:
+        >>> from pymss import resolve_model
+        >>> resolved = resolve_model("bs_roformer_voc_hyperacev2", require_exists=False)
+        >>> resolved["model_type"]
+        'bs_roformer'
+
+    Example:
+        >>> resolved = resolve_model("bs_roformer_voc_hyperacev2", model_dir="models")
+        >>> resolved["model_path"].endswith(".ckpt") or resolved["model_path"].endswith(".pth")
+        True"""
+    entry = get_model_entry(model_name)
+    if require_supported and not entry.supported:
+        reason = entry.unsupported_reason or "unsupported"
+        raise ValueError(f"Model {entry.name} cannot be used for inference yet: {reason}")
+
+    model_path = model_path_for(entry, model_dir)
+    config_path = config_path_for(entry, model_dir)
+    missing = []
+    if require_exists and not model_path.is_file():
+        missing.append(str(model_path))
+    if require_exists and config_path is not None and not config_path.is_file():
+        missing.append(str(config_path))
+    for path in auxiliary_paths_for(entry, model_dir):
+        if require_exists and not path.is_file():
+            missing.append(str(path))
+    if missing:
+        raise FileNotFoundError("Missing model file(s): " + ", ".join(missing))
+
+    return {
+        "entry": entry,
+        "model_type": entry.model_type,
+        "model_path": str(model_path),
+        "config_path": str(config_path) if config_path else None,
+    }
+
+
+def create_separator(model_name, model_dir=None, **separator_kwargs):
+    """Create ``MSSeparator`` from a catalog model name.
+
+    This is a convenience wrapper around ``resolve_model(...)`` followed by
+    ``MSSeparator(...)``. It expects the model files to already exist locally;
+    call ``download_model(...)`` first or use ``MSSeparator.from_model_name`` if
+    you want optional downloading in one step.
+
+    Args:
+        model_name (str): Model name, stem, or alias from the pymss catalog.
+        model_dir (str | os.PathLike | None, optional): Local model cache
+            directory. Defaults to None.
+        **separator_kwargs: Keyword arguments forwarded to ``MSSeparator``,
+            such as ``device``, ``device_ids``, ``output_format``,
+            ``store_dirs``, ``save_as_folder``, ``audio_params``, ``logger``,
+            ``debug``, ``progress_callback``, and ``inference_params``.
+
+    Returns:
+        MSSeparator: Loaded separator instance ready for inference.
+
+    Raises:
+        FileNotFoundError: If required model files are not present locally.
+
+    Example:
+        >>> from pymss import create_separator
+        >>> separator = create_separator(
+        ...     "bs_roformer_voc_hyperacev2",
+        ...     model_dir="models",
+        ...     output_format="wav",
+        ...     inference_params={"normalize": True},
+        ... )
+        >>> separator.process_folder("song.wav")
+
+    Example:
+        >>> separator = create_separator(
+        ...     "some_six_stem_model",
+        ...     store_dirs={"vocals": "out/vocals", "drums": "out/drums"},
+        ... )"""
+    from .separator import MSSeparator
+
+    resolved = resolve_model(model_name, model_dir=model_dir, require_supported=True, require_exists=True)
+    return MSSeparator(
+        model_type=resolved["model_type"],
+        model_path=resolved["model_path"],
+        config_path=resolved["config_path"],
+        **separator_kwargs,
+    )

tools/pymss/modules/__init__.py

@@ -0,0 +1,62 @@
+from ._core_shims import alias_submodules
+
+alias_submodules(
+    __name__,
+    "pymss_core.modules",
+    (
+        "apollo_mlx",
+        "bandit_mlx",
+        "demucs4ht",
+        "demucs_local",
+        "demucs_mlx",
+        "legacy_demucs",
+        "mdx23c_mlx",
+        "mdx23c_tfc_tdf_v3",
+        "mlx_utils",
+        "scnet_mlx",
+        "spectrogram",
+    ),
+)
+
+alias_submodules(
+    __name__,
+    "pymss_core.modules",
+    (
+        "bandit",
+        "bandit.bandsplit",
+        "bandit.core",
+        "bandit.core.model",
+        "bandit.core.model._spectral",
+        "bandit.core.model.bsrnn",
+        "bandit.core.model.bsrnn.bandsplit",
+        "bandit.core.model.bsrnn.core",
+        "bandit.core.model.bsrnn.maskestim",
+        "bandit.core.model.bsrnn.tfmodel",
+        "bandit.core.model.bsrnn.utils",
+        "bandit.core.model.bsrnn.wrapper",
+        "bandit.maskestim",
+        "bandit.tfmodel",
+        "bandit_v2",
+        "bandit_v2.bandit",
+        "bandit_v2.bandsplit",
+        "bandit_v2.maskestim",
+        "bandit_v2.tfmodel",
+        "bandit_v2.utils",
+        "bs_roformer",
+        "bs_roformer.attend",
+        "bs_roformer.bands",
+        "bs_roformer.bs_roformer",
+        "bs_roformer.bs_roformer_hyperace",
+        "bs_roformer.common",
+        "bs_roformer.hyperace_segm",
+        "bs_roformer.mel_band_roformer",
+        "bs_roformer.mlx_attention",
+        "bs_roformer.mlx_roformer",
+        "bs_roformer.transformer",
+        "look2hear",
+        "look2hear.apollo",
+        "scnet",
+        "scnet.scnet",
+        "scnet.separation",
+    ),
+)

tools/pymss/modules/_core_shims.py

@@ -0,0 +1,20 @@
+from importlib import import_module
+import sys
+
+_LOCAL_MODULE_PREFIX = "pymss.modules."
+_CORE_MODULE_PREFIX = "pymss_core.modules."
+
+
+def alias_module(local_name, core_name):
+    if not local_name.startswith(_LOCAL_MODULE_PREFIX):
+        raise ValueError(f"invalid local module alias: {local_name}")
+    if not core_name.startswith(_CORE_MODULE_PREFIX):
+        raise ValueError(f"invalid core module alias: {core_name}")
+    module = import_module(core_name)
+    sys.modules[local_name] = module
+    return module
+
+
+def alias_submodules(local_package, core_package, names):
+    for name in names:
+        alias_module(f"{local_package}.{name}", f"{core_package}.{name}")

tools/pymss/modules/vocal_remover/__init__.py

@@ -0,0 +1,20 @@
+from pymss_core.modules.vocal_remover import (
+    BaseASPPNet,
+    BaseNet,
+    CascadedASPPNet,
+    CascadedNet,
+    ModelParameters,
+    determine_model_capacity,
+)
+
+from .vr_separator import VRSeparator
+
+__all__ = (
+    "BaseASPPNet",
+    "BaseNet",
+    "CascadedASPPNet",
+    "CascadedNet",
+    "ModelParameters",
+    "VRSeparator",
+    "determine_model_capacity",
+)

tools/pymss/modules/vocal_remover/common_separator.py

@@ -0,0 +1,53 @@
+class CommonSeparator:
+    VOCAL_STEM = "Vocals"
+    OTHER_STEM = "Other"
+    BASS_STEM = "Bass"
+    DRUM_STEM = "Drums"
+    GUITAR_STEM = "Guitar"
+    PIANO_STEM = "Piano"
+    SYNTH_STEM = "Synthesizer"
+    STRINGS_STEM = "Strings"
+    WOODWINDS_STEM = "Woodwinds"
+    BRASS_STEM = "Brass"
+    WIND_INST_STEM = "Wind Inst"
+
+    NON_ACCOM_STEMS = (
+        VOCAL_STEM,
+        OTHER_STEM,
+        BASS_STEM,
+        DRUM_STEM,
+        GUITAR_STEM,
+        PIANO_STEM,
+        SYNTH_STEM,
+        STRINGS_STEM,
+        WOODWINDS_STEM,
+        BRASS_STEM,
+        WIND_INST_STEM,
+    )
+
+    def __init__(self, config):
+        self.logger = config.get("logger")
+        self.debug = config.get("debug")
+        self.torch_device = config.get("torch_device")
+        self.torch_device_cpu = config.get("torch_device_cpu")
+        self.torch_device_mps = config.get("torch_device_mps")
+        self.model_name = config.get("model_name")
+        self.model_path = config.get("model_path")
+        self.model_data = config.get("model_data")
+        self.sample_rate = config.get("sample_rate")
+        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")
+        self.is_karaoke = self.model_data.get("is_karaoke", False)
+        self.is_bv_model = self.model_data.get("is_bv_model", False)
+        self.bv_model_rebalance = self.model_data.get("is_bv_model_rebalanced", 0)
+
+        self.primary_source = None
+        self.secondary_source = None
+
+        self.logger.info(f"VR params: model_name={self.model_name}, model_path={self.model_path}")
+        self.logger.info(f"VR params: primary_stem={self.primary_stem_name}, secondary_stem={self.secondary_stem_name}")
+        self.logger.debug(
+            f"VR params: is_karaoke={self.is_karaoke}, is_bv_model={self.is_bv_model}, bv_model_rebalance={self.bv_model_rebalance}"
+        )

tools/pymss/modules/vocal_remover/uvr_lib_v5/__init__.py

@@ -0,0 +1,14 @@
+from pymss.modules._core_shims import alias_submodules
+
+alias_submodules(
+    __name__,
+    "pymss_core.modules.vocal_remover.uvr_lib_v5",
+    (
+        "vr_network",
+        "vr_network.layers",
+        "vr_network.layers_new",
+        "vr_network.model_param_init",
+        "vr_network.nets",
+        "vr_network.nets_new",
+    ),
+)

tools/pymss/modules/vocal_remover/uvr_lib_v5/spec_utils.py

@@ -0,0 +1,388 @@
+import math
+import platform
+import traceback
+
+import librosa
+import numpy as np
+import torch
+
+
+ARM = "arm"
+wav_resolution = (
+    "polyphase" if platform.system() == "Darwin" and (platform.processor() == ARM or ARM in platform.platform()) else "soxr_hq"
+)
+
+
+_HANN_WINDOW_CACHE = {}
+_FILTER_MASK_CACHE = {}
+
+
+def _hann_window(n_fft, dtype, device):
+    torch_device = torch.device(device)
+    key = (int(n_fft), dtype, torch_device.type, torch_device.index)
+    window = _HANN_WINDOW_CACHE.get(key)
+    if window is None:
+        window = torch.hann_window(n_fft, dtype=dtype, device=torch_device)
+        _HANN_WINDOW_CACHE[key] = window
+    return window
+
+
+def resample_audio(wave, orig_sr, target_sr, res_type=None):
+    orig_sr = int(orig_sr)
+    target_sr = int(target_sr)
+    if orig_sr == target_sr:
+        return np.asfortranarray(wave)
+
+    last_error = None
+    for candidate in dict.fromkeys(([res_type] if res_type else []) + ["soxr_hq", "polyphase"]):
+        try:
+            return librosa.resample(wave, orig_sr=orig_sr, target_sr=target_sr, res_type=candidate)
+        except (ImportError, ModuleNotFoundError) as exc:
+            last_error = exc
+
+    try:
+        return _linear_resample(wave, orig_sr, target_sr)
+    except Exception:
+        if last_error is not None:
+            raise last_error
+        raise
+
+
+def _linear_resample(wave, orig_sr, target_sr):
+    wave = np.asarray(wave)
+    original_length = wave.shape[-1]
+    target_length = max(1, int(round(original_length * target_sr / orig_sr)))
+    if original_length == target_length:
+        return np.asfortranarray(wave)
+
+    old_x, new_x = np.linspace(0.0, 1.0, original_length, endpoint=False), np.linspace(0.0, 1.0, target_length, endpoint=False)
+    if wave.ndim == 1:
+        return np.asfortranarray(np.interp(new_x, old_x, wave).astype(wave.dtype, copy=False))
+    return np.asfortranarray(
+        np.stack(
+            [
+                np.interp(new_x, old_x, channel).astype(wave.dtype, copy=False)
+                for channel in wave.reshape((-1, original_length))
+            ],
+            axis=0,
+        ).reshape(wave.shape[:-1] + (target_length,))
+    )
+
+
+def crop_center(h1, h2):
+    h1_time, h2_time = h1.size(3), h2.size(3)
+    if h1_time == h2_time:
+        return h1
+    if h1_time < h2_time:
+        raise ValueError("h1_shape[3] must be greater than h2_shape[3]")
+    start = (h1_time - h2_time) // 2
+    return h1[:, :, :, start : start + h2_time]
+
+
+def preprocess(x_spec):
+    return np.abs(x_spec), np.angle(x_spec)
+
+
+def make_padding(width, cropsize, offset):
+    roi_size = cropsize - offset * 2 or cropsize
+    return offset, roi_size - (width % roi_size) + offset, roi_size
+
+
+def merge_artifacts(y_mask, thres=0.01, min_range=64, fade_size=32):
+    mask = y_mask
+    try:
+        if min_range < fade_size * 2:
+            raise ValueError("min_range must be >= fade_size * 2")
+
+        idx = np.where(y_mask.min(axis=(0, 1)) > thres)[0]
+        if len(idx) == 0:
+            return mask
+        start_idx = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0])
+        end_idx = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1])
+        artifact_idx = np.where(end_idx - start_idx > min_range)[0]
+        weight = np.zeros_like(y_mask)
+        if len(artifact_idx) > 0:
+            start_idx = start_idx[artifact_idx]
+            end_idx = end_idx[artifact_idx]
+            old_e = None
+            for s, e in zip(start_idx, end_idx):
+                if old_e is not None and s - old_e < fade_size:
+                    s = old_e - fade_size * 2
+                if s != 0:
+                    weight[:, :, s : s + fade_size] = np.linspace(0, 1, fade_size)
+                else:
+                    s -= fade_size
+                if e != y_mask.shape[2]:
+                    weight[:, :, e - fade_size : e] = np.linspace(1, 0, fade_size)
+                else:
+                    e += fade_size
+                weight[:, :, s + fade_size : e - fade_size] = 1
+                old_e = e
+
+        y_mask += weight * (1 - y_mask)
+        mask = y_mask
+    except Exception as exc:
+        print(f'Post Process Failed: {type(exc).__name__}: "{exc}"\n{"".join(traceback.format_tb(exc.__traceback__))}"')
+    return mask
+
+
+def convert_channels(spec, mp, band):
+    mode = mp.param["band"][band].get("convert_channels")
+    if mode == "mid_side_c":
+        return np.asfortranarray([np.add(spec[0], spec[1] * 0.25), np.subtract(spec[1], spec[0] * 0.25)])
+    if mode == "mid_side":
+        return np.asfortranarray([np.add(spec[0], spec[1]) / 2, np.subtract(spec[0], spec[1])])
+    if mode == "stereo_n":
+        return np.asfortranarray([np.add(spec[0], spec[1] * 0.25) / 0.9375, np.add(spec[1], spec[0] * 0.25) / 0.9375])
+    return spec
+
+
+def combine_spectrograms(specs, mp, is_v51_model=False):
+    length = min(specs[i].shape[2] for i in specs)
+    spec_c = np.zeros((2, mp.param["bins"] + 1, length), dtype=np.complex64)
+    offset = 0
+    bands_n = len(mp.param["band"])
+    pre_start, pre_stop = mp.param["pre_filter_start"], mp.param["pre_filter_stop"]
+
+    for d in range(1, bands_n + 1):
+        band = mp.param["band"][d]
+        height = band["crop_stop"] - band["crop_start"]
+        spec_c[:, offset : offset + height, :length] = specs[d][:, band["crop_start"] : band["crop_stop"], :length]
+        offset += height
+
+    if offset > mp.param["bins"]:
+        raise ValueError("Too much bins")
+
+    if pre_start > 0:
+        if is_v51_model:
+            spec_c *= get_lp_filter_mask(spec_c.shape[1], pre_start, pre_stop)
+        elif bands_n == 1:
+            spec_c = fft_lp_filter(spec_c, pre_start, pre_stop)
+        else:
+            gain_prev = 1
+            for b in range(pre_start + 1, pre_stop):
+                gain = math.pow(10, -(b - pre_start) * (3.5 - gain_prev) / 20.0)
+                gain_prev = gain
+                spec_c[:, b, :] *= gain
+
+    return np.asfortranarray(spec_c)
+
+
+def wave_to_spectrogram(wave, hop_length, n_fft, mp, band, is_v51_model=False, torch_device=None):
+    if wave.ndim == 1:
+        wave = np.asfortranarray([wave, wave])
+
+    left, right = (np.asfortranarray(channel) for channel in wave[:2])
+    if not is_v51_model:
+        if mp.param["reverse"]:
+            left, right = np.flip(left), np.flip(right)
+        elif mp.param["mid_side"]:
+            left, right = np.asfortranarray(np.add(wave[0], wave[1]) / 2), np.asfortranarray(np.subtract(wave[0], wave[1]))
+        elif mp.param["mid_side_b2"]:
+            left, right = (
+                np.asfortranarray(np.add(wave[1], wave[0] * 0.5)),
+                np.asfortranarray(np.subtract(wave[0], wave[1] * 0.5)),
+            )
+
+    spec = _torch_stft(np.asfortranarray([left, right]), n_fft, hop_length, torch_device)
+    if spec is None:
+        spec = np.asfortranarray(
+            [
+                librosa.stft(left, n_fft=n_fft, hop_length=hop_length),
+                librosa.stft(right, n_fft=n_fft, hop_length=hop_length),
+            ]
+        )
+    return convert_channels(spec, mp, band) if is_v51_model else spec
+
+
+def _torch_stft(wave, n_fft, hop_length, device):
+    if device is None or torch.device(device).type != "cuda":
+        return None
+    wave_t = torch.from_numpy(np.ascontiguousarray(wave)).to(device)
+    window = _hann_window(n_fft, wave_t.dtype, device)
+    spec = torch.stft(
+        wave_t,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        window=window,
+        center=True,
+        pad_mode="constant",
+        return_complex=True,
+    )
+    return np.asfortranarray(spec.cpu().numpy())
+
+
+def _torch_istft(spec, hop_length, device):
+    if device is None or torch.device(device).type != "cuda":
+        return None
+    n_fft = (spec.shape[1] - 1) * 2
+    spec_t = torch.from_numpy(np.ascontiguousarray(spec)).to(device)
+    window = _hann_window(n_fft, spec_t.real.dtype, device)
+    wave = torch.istft(spec_t, n_fft=n_fft, hop_length=hop_length, window=window, center=True, return_complex=False)
+    return np.asfortranarray(wave.cpu().numpy())
+
+
+def spectrogram_to_wave(spec, hop_length=1024, mp=None, band=0, is_v51_model=True, torch_device=None):
+    wave = _torch_istft(spec, hop_length, torch_device)
+    if wave is None:
+        left = librosa.istft(np.asfortranarray(spec[0]), hop_length=hop_length, dtype=np.float64)
+        right = librosa.istft(np.asfortranarray(spec[1]), hop_length=hop_length, dtype=np.float64)
+    else:
+        left, right = wave[0], wave[1]
+
+    if is_v51_model:
+        mode = mp.param["band"][band].get("convert_channels")
+        if mode == "mid_side_c":
+            return np.asfortranarray([np.subtract(left / 1.0625, right / 4.25), np.add(right / 1.0625, left / 4.25)])
+        if mode == "mid_side":
+            return np.asfortranarray([np.add(left, right / 2), np.subtract(left, right / 2)])
+        if mode == "stereo_n":
+            return np.asfortranarray([np.subtract(left, right * 0.25), np.subtract(right, left * 0.25)])
+    else:
+        if mp.param["reverse"]:
+            return np.asfortranarray([np.flip(left), np.flip(right)])
+        if mp.param["mid_side"]:
+            return np.asfortranarray([np.add(left, right / 2), np.subtract(left, right / 2)])
+        if mp.param["mid_side_b2"]:
+            return np.asfortranarray([np.add(right / 1.25, 0.4 * left), np.subtract(left / 1.25, 0.4 * right)])
+
+    return np.asfortranarray([left, right])
+
+
+def cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None, is_v51_model=False, torch_device=None):
+    spec_m = np.where(np.isnan(spec_m), 0, spec_m)
+    extra_bins_h = None if extra_bins_h is None else int(extra_bins_h)
+    extra_bins = None if extra_bins is None else np.where(np.isnan(extra_bins), 0, extra_bins)
+
+    bands_n = len(mp.param["band"])
+    offset = 0
+    wave = None
+
+    for d in range(1, bands_n + 1):
+        bp = mp.param["band"][d]
+        spec_s = np.zeros((2, bp["n_fft"] // 2 + 1, spec_m.shape[2]), dtype=np.result_type(spec_m.dtype, np.complex64))
+        height = bp["crop_stop"] - bp["crop_start"]
+        spec_s[:, bp["crop_start"] : bp["crop_stop"], :] = spec_m[:, offset : offset + height, :]
+        offset += height
+
+        if d == bands_n:
+            if extra_bins_h is not None:
+                spec_s[:, bp["n_fft"] // 2 - extra_bins_h : bp["n_fft"] // 2, :] = extra_bins[:, :extra_bins_h, :]
+            if bp["hpf_start"] > 0:
+                spec_s = (
+                    spec_s * get_hp_filter_mask(spec_s.shape[1], bp["hpf_start"], bp["hpf_stop"] - 1)
+                    if is_v51_model
+                    else fft_hp_filter(spec_s, bp["hpf_start"], bp["hpf_stop"] - 1)
+                )
+            band_wave = spectrogram_to_wave(spec_s, bp["hl"], mp, d, is_v51_model, torch_device=torch_device)
+            wave = band_wave if wave is None else np.add(wave, band_wave)
+        else:
+            sr = mp.param["band"][d + 1]["sr"]
+            if d == 1:
+                spec_s = (
+                    spec_s * get_lp_filter_mask(spec_s.shape[1], bp["lpf_start"], bp["lpf_stop"])
+                    if is_v51_model
+                    else fft_lp_filter(spec_s, bp["lpf_start"], bp["lpf_stop"])
+                )
+                wave = resample_audio(
+                    spectrogram_to_wave(spec_s, bp["hl"], mp, d, is_v51_model, torch_device=torch_device),
+                    orig_sr=bp["sr"],
+                    target_sr=sr,
+                    res_type=wav_resolution,
+                )
+            else:
+                if is_v51_model:
+                    spec_s *= get_hp_filter_mask(spec_s.shape[1], bp["hpf_start"], bp["hpf_stop"] - 1)
+                    spec_s *= get_lp_filter_mask(spec_s.shape[1], bp["lpf_start"], bp["lpf_stop"])
+                else:
+                    spec_s = fft_hp_filter(spec_s, bp["hpf_start"], bp["hpf_stop"] - 1)
+                    spec_s = fft_lp_filter(spec_s, bp["lpf_start"], bp["lpf_stop"])
+                wave = resample_audio(
+                    np.add(wave, spectrogram_to_wave(spec_s, bp["hl"], mp, d, is_v51_model, torch_device=torch_device)),
+                    orig_sr=bp["sr"],
+                    target_sr=sr,
+                    res_type=wav_resolution,
+                )
+
+    return wave
+
+
+def _get_filter_mask(kind, n_bins, bin_start, bin_stop):
+    key = (kind, int(n_bins), int(bin_start), int(bin_stop))
+    mask = _FILTER_MASK_CACHE.get(key)
+    if mask is None:
+        mask = np.concatenate(
+            (
+                [
+                    np.ones((bin_start - 1, 1)),
+                    np.linspace(1, 0, bin_stop - bin_start + 1)[:, None],
+                    np.zeros((n_bins - bin_stop, 1)),
+                ]
+                if kind == "lp"
+                else [
+                    np.zeros((bin_stop + 1, 1)),
+                    np.linspace(0, 1, 1 + bin_start - bin_stop)[:, None],
+                    np.ones((n_bins - bin_start - 2, 1)),
+                ]
+            ),
+            axis=0,
+        )
+        _FILTER_MASK_CACHE[key] = mask
+    return mask
+
+
+def get_lp_filter_mask(n_bins, bin_start, bin_stop):
+    return _get_filter_mask("lp", n_bins, bin_start, bin_stop)
+
+
+def get_hp_filter_mask(n_bins, bin_start, bin_stop):
+    return _get_filter_mask("hp", n_bins, bin_start, bin_stop)
+
+
+def fft_lp_filter(spec, bin_start, bin_stop):
+    gain = 1.0
+    for b in range(bin_start, bin_stop):
+        gain -= 1 / (bin_stop - bin_start)
+        spec[:, b, :] = gain * spec[:, b, :]
+    spec[:, bin_stop:, :] *= 0
+    return spec
+
+
+def fft_hp_filter(spec, bin_start, bin_stop):
+    gain = 1.0
+    for b in range(bin_start, bin_stop, -1):
+        gain -= 1 / (bin_start - bin_stop)
+        spec[:, b, :] = gain * spec[:, b, :]
+    spec[:, 0 : bin_stop + 1, :] *= 0
+    return spec
+
+
+def mirroring(mode, spec_m, input_high_end, mp):
+    if mode not in ("mirroring", "mirroring2"):
+        return input_high_end
+    mirror = np.flip(
+        np.abs(spec_m[:, mp.param["pre_filter_start"] - 10 - input_high_end.shape[1] : mp.param["pre_filter_start"] - 10, :]), 1
+    )
+    mirror = (
+        mirror * np.exp(1.0j * np.angle(input_high_end)) if mode == "mirroring" else np.multiply(mirror, input_high_end * 1.7)
+    )
+    return np.where(np.abs(input_high_end) <= np.abs(mirror), input_high_end, mirror)
+
+
+def adjust_aggr(mask, is_non_accom_stem, aggressiveness):
+    aggr = aggressiveness["value"] * 2
+    if aggr != 0:
+        if is_non_accom_stem:
+            aggr = 1 - aggr
+        if np.any(aggr > 10) or np.any(aggr < -10):
+            print(f"Warning: Extreme aggressiveness values detected: {aggr}")
+
+        aggr = np.array([aggr, aggr])
+        if (correction := aggressiveness["aggr_correction"]) is not None:
+            aggr[0] += correction["left"]
+            aggr[1] += correction["right"]
+
+        split_bin = aggressiveness["split_bin"]
+        mask[:, :split_bin] = np.power(mask[:, :split_bin], 1 + aggr[:, None, None] / 3)
+        mask[:, split_bin:] = np.power(mask[:, split_bin:], 1 + aggr[:, None, None])
+    return mask

tools/pymss/modules/vocal_remover/vr_mlx.py

@@ -0,0 +1,407 @@
+import torch
+
+from ..bs_roformer.mlx_attention import _linear, _mlx_dtype, _torch_to_mlx_array
+from pymss_core.modules.vocal_remover.uvr_lib_v5.vr_network import layers, layers_new, nets, nets_new
+
+
+def _mlx_param(module, name, tensor, dtype):
+    cache = getattr(module, "_pymss_mlx_full_param_cache", None)
+    if cache is None:
+        cache = {}
+        module._pymss_mlx_full_param_cache = cache
+    key = (name, tensor.data_ptr(), tensor._version, tuple(tensor.shape), dtype)
+    cached = cache.get(name)
+    if cached is not None and cached[0] == key:
+        return cached[1]
+    value = _torch_to_mlx_array(tensor, dtype)
+    cache[name] = (key, value)
+    return value
+
+
+def _conv_padding(conv):
+    padding = conv.padding
+    if isinstance(padding, tuple):
+        return padding
+    return padding, padding
+
+
+def _conv2d_nchw(conv, x, dtype):
+    import mlx.core as mx
+
+    weight = mx.transpose(_mlx_param(conv, "weight", conv.weight, dtype), (0, 2, 3, 1))
+    y = mx.conv2d(
+        mx.transpose(x, (0, 2, 3, 1)),
+        weight,
+        stride=conv.stride,
+        padding=_conv_padding(conv),
+        dilation=conv.dilation,
+        groups=conv.groups,
+    )
+    if conv.bias is not None:
+        y = y + _mlx_param(conv, "bias", conv.bias, dtype)
+    return mx.transpose(y, (0, 3, 1, 2))
+
+
+def _batch_norm2d(module, x, dtype):
+    import mlx.core as mx
+
+    if module.training:
+        raise TypeError("MLX VR BatchNorm2d supports eval mode only")
+    y = x.astype(mx.float32)
+    mean = _torch_to_mlx_array(module.running_mean, torch.float32).reshape(1, -1, 1, 1)
+    var = _torch_to_mlx_array(module.running_var, torch.float32).reshape(1, -1, 1, 1)
+    y = (y - mean) * mx.rsqrt(var + module.eps)
+    if module.affine:
+        weight = _mlx_param(module, "weight", module.weight, dtype).reshape(1, -1, 1, 1)
+        bias = _mlx_param(module, "bias", module.bias, dtype).reshape(1, -1, 1, 1)
+        y = y.astype(x.dtype) * weight + bias
+    return y.astype(x.dtype)
+
+
+def _batch_norm1d(module, x, dtype):
+    import mlx.core as mx
+
+    if module.training:
+        raise TypeError("MLX VR BatchNorm1d supports eval mode only")
+    y = x.astype(mx.float32)
+    mean = _torch_to_mlx_array(module.running_mean, torch.float32).reshape(1, -1)
+    var = _torch_to_mlx_array(module.running_var, torch.float32).reshape(1, -1)
+    y = (y - mean) * mx.rsqrt(var + module.eps)
+    if module.affine:
+        y = y.astype(x.dtype) * _mlx_param(module, "weight", module.weight, dtype).reshape(1, -1)
+        y = y + _mlx_param(module, "bias", module.bias, dtype).reshape(1, -1)
+    return y.astype(x.dtype)
+
+
+def _activation(module, x):
+    import mlx.core as mx
+
+    if isinstance(module, torch.nn.ReLU):
+        return mx.maximum(x, 0)
+    if isinstance(module, torch.nn.LeakyReLU):
+        return mx.maximum(x, 0) + module.negative_slope * mx.minimum(x, 0)
+    if isinstance(module, torch.nn.Sigmoid):
+        return mx.sigmoid(x)
+    if isinstance(module, (torch.nn.Dropout, torch.nn.Dropout2d, torch.nn.Identity)):
+        return x
+    raise TypeError(f"unsupported VR activation for MLX full backend: {type(module).__name__}")
+
+
+def _resize_positions_align_corners(in_size, out_size):
+    import mlx.core as mx
+
+    if out_size == 1:
+        pos = mx.zeros((1,), dtype=mx.float32)
+    else:
+        pos = mx.arange(out_size, dtype=mx.float32) * ((in_size - 1) / (out_size - 1))
+    lower = mx.floor(pos)
+    upper = lower + 1
+    weight = pos - lower
+    lower = mx.clip(lower, 0, in_size - 1).astype(mx.int32)
+    upper = mx.clip(upper, 0, in_size - 1).astype(mx.int32)
+    return lower, upper, weight
+
+
+def _resize_bilinear_nchw(x, size=None, scale_factor=None):
+    import mlx.core as mx
+
+    if size is None:
+        out_h = int(x.shape[2] * scale_factor)
+        out_w = int(x.shape[3] * scale_factor)
+    else:
+        out_h, out_w = int(size[0]), int(size[1])
+    in_h, in_w = x.shape[2], x.shape[3]
+    if in_h == out_h and in_w == out_w:
+        return x
+    y0, y1, wy = _resize_positions_align_corners(in_h, out_h)
+    x0, x1, wx = _resize_positions_align_corners(in_w, out_w)
+    v00 = mx.take(mx.take(x, y0, axis=2), x0, axis=3)
+    v01 = mx.take(mx.take(x, y0, axis=2), x1, axis=3)
+    v10 = mx.take(mx.take(x, y1, axis=2), x0, axis=3)
+    v11 = mx.take(mx.take(x, y1, axis=2), x1, axis=3)
+    wy = wy.reshape(1, 1, out_h, 1)
+    wx = wx.reshape(1, 1, 1, out_w)
+    return v00 * (1 - wy) * (1 - wx) + v01 * (1 - wy) * wx + v10 * wy * (1 - wx) + v11 * wy * wx
+
+
+def _crop_center(skip, target):
+    h, w = target.shape[2], target.shape[3]
+    dh = (skip.shape[2] - h) // 2
+    dw = (skip.shape[3] - w) // 2
+    return skip[:, :, dh : dh + h, dw : dw + w]
+
+
+def _adaptive_avg_pool_1_none(x):
+    import mlx.core as mx
+
+    return mx.mean(x, axis=2, keepdims=True)
+
+
+def _replicate_pad_freq_bottom(x, pad):
+    import mlx.core as mx
+
+    if pad <= 0:
+        return x
+    last = mx.broadcast_to(x[:, :, -1:, :], (x.shape[0], x.shape[1], pad, x.shape[3]))
+    return mx.concatenate((x, last), axis=2)
+
+
+def _seq(module, x, dtype):
+    for child in module:
+        x = _module_forward(child, x, dtype)
+    return x
+
+
+def _module_forward(module, x, dtype):
+    if isinstance(module, torch.nn.Sequential):
+        return _seq(module, x, dtype)
+    if isinstance(module, (layers.Conv2DBNActiv, layers.SeperableConv2DBNActiv, layers_new.Conv2DBNActiv)):
+        return _seq(module.conv, x, dtype)
+    if isinstance(module, torch.nn.Conv2d):
+        return _conv2d_nchw(module, x, dtype)
+    if isinstance(module, torch.nn.BatchNorm2d):
+        return _batch_norm2d(module, x, dtype)
+    if isinstance(module, torch.nn.BatchNorm1d):
+        return _batch_norm1d(module, x, dtype)
+    if isinstance(module, torch.nn.Linear):
+        return _linear(
+            x,
+            _mlx_param(module, "weight", module.weight, dtype),
+            None if module.bias is None else _mlx_param(module, "bias", module.bias, dtype),
+        )
+    if isinstance(module, torch.nn.AdaptiveAvgPool2d):
+        if module.output_size != (1, None):
+            raise TypeError(f"unsupported VR AdaptiveAvgPool2d output_size: {module.output_size}")
+        return _adaptive_avg_pool_1_none(x)
+    if isinstance(
+        module, (torch.nn.ReLU, torch.nn.LeakyReLU, torch.nn.Sigmoid, torch.nn.Dropout, torch.nn.Dropout2d, torch.nn.Identity)
+    ):
+        return _activation(module, x)
+    if isinstance(module, layers.ASPPModule):
+        return _old_aspp(module, x, dtype)
+    if isinstance(module, layers.Decoder):
+        return _old_decoder(module, x, None, dtype)
+    if isinstance(module, layers_new.ASPPModule):
+        return _new_aspp(module, x, dtype)
+    if isinstance(module, layers_new.LSTMModule):
+        return _new_lstm_module(module, x, dtype)
+    if isinstance(module, layers_new.Decoder):
+        return _new_decoder(module, x, None, dtype)
+    if isinstance(module, nets.BaseASPPNet):
+        return _old_base_aspp_net(module, x, dtype)
+    if isinstance(module, nets_new.BaseNet):
+        return _new_base_net(module, x, dtype)
+    raise TypeError(f"unsupported VR layer for MLX full backend: {type(module).__name__}")
+
+
+def _old_encoder(module, x, dtype):
+    skip = _module_forward(module.conv1, x, dtype)
+    return _module_forward(module.conv2, skip, dtype), skip
+
+
+def _old_decoder(module, x, skip, dtype):
+    import mlx.core as mx
+
+    x = _resize_bilinear_nchw(x, scale_factor=2)
+    if skip is not None:
+        x = mx.concatenate((x, _crop_center(skip, x)), axis=1)
+    x = _module_forward(module.conv, x, dtype)
+    return x if module.dropout is None else x
+
+
+def _old_aspp(module, x, dtype):
+    import mlx.core as mx
+
+    h, w = x.shape[2], x.shape[3]
+    features = [
+        _resize_bilinear_nchw(_module_forward(module.conv1, x, dtype), size=(h, w)),
+        _module_forward(module.conv2, x, dtype),
+        _module_forward(module.conv3, x, dtype),
+        _module_forward(module.conv4, x, dtype),
+        _module_forward(module.conv5, x, dtype),
+    ]
+    if module.nn_architecture in module.six_layer:
+        features.append(_module_forward(module.conv6, x, dtype))
+    elif module.nn_architecture in module.seven_layer:
+        features.extend((_module_forward(module.conv6, x, dtype), _module_forward(module.conv7, x, dtype)))
+    return _module_forward(module.bottleneck, mx.concatenate(features, axis=1), dtype)
+
+
+def _old_base_aspp_net(module, x, dtype):
+    x, skip1 = _old_encoder(module.enc1, x, dtype)
+    x, skip2 = _old_encoder(module.enc2, x, dtype)
+    x, skip3 = _old_encoder(module.enc3, x, dtype)
+    x, skip4 = _old_encoder(module.enc4, x, dtype)
+
+    if module.nn_architecture == 129605:
+        x, skip5 = _old_encoder(module.enc5, x, dtype)
+        x = _old_decoder(module.dec5, _old_aspp(module.aspp, x, dtype), skip5, dtype)
+    else:
+        x = _old_aspp(module.aspp, x, dtype)
+
+    x = _old_decoder(module.dec4, x, skip4, dtype)
+    x = _old_decoder(module.dec3, x, skip3, dtype)
+    x = _old_decoder(module.dec2, x, skip2, dtype)
+    return _old_decoder(module.dec1, x, skip1, dtype)
+
+
+def _old_cascaded_aspp_net(module, x, dtype):
+    import mlx.core as mx
+
+    x = x[:, :, : module.max_bin]
+    bandwidth = x.shape[2] // 2
+    aux1 = mx.concatenate(
+        (
+            _old_base_aspp_net(module.stg1_low_band_net, x[:, :, :bandwidth], dtype),
+            _old_base_aspp_net(module.stg1_high_band_net, x[:, :, bandwidth:], dtype),
+        ),
+        axis=2,
+    )
+
+    hidden = mx.concatenate((x, aux1), axis=1)
+    aux2 = _old_base_aspp_net(module.stg2_full_band_net, _module_forward(module.stg2_bridge, hidden, dtype), dtype)
+
+    hidden = mx.concatenate((x, aux1, aux2), axis=1)
+    mask = mx.sigmoid(
+        _conv2d_nchw(
+            module.out,
+            _old_base_aspp_net(module.stg3_full_band_net, _module_forward(module.stg3_bridge, hidden, dtype), dtype),
+            dtype,
+        )
+    )
+    return _replicate_pad_freq_bottom(mask, module.output_bin - mask.shape[2])
+
+
+def _new_encoder(module, x, dtype):
+    return _module_forward(module.conv2, _module_forward(module.conv1, x, dtype), dtype)
+
+
+def _new_decoder(module, x, skip, dtype):
+    import mlx.core as mx
+
+    x = _resize_bilinear_nchw(x, scale_factor=2)
+    if skip is not None:
+        x = mx.concatenate((x, _crop_center(skip, x)), axis=1)
+    x = _module_forward(module.conv1, x, dtype)
+    return x if module.dropout is None else x
+
+
+def _new_aspp(module, x, dtype):
+    import mlx.core as mx
+
+    h, w = x.shape[2], x.shape[3]
+    out = mx.concatenate(
+        (
+            _resize_bilinear_nchw(_module_forward(module.conv1, x, dtype), size=(h, w)),
+            _module_forward(module.conv2, x, dtype),
+            _module_forward(module.conv3, x, dtype),
+            _module_forward(module.conv4, x, dtype),
+            _module_forward(module.conv5, x, dtype),
+        ),
+        axis=1,
+    )
+    out = _module_forward(module.bottleneck, out, dtype)
+    return out if module.dropout is None else out
+
+
+def _lstm_forward(rnn, x, dtype):
+    import mlx.core as mx
+
+    if rnn.num_layers != 1 or rnn.batch_first:
+        raise TypeError("MLX VR LSTM supports one-layer non-batch-first LSTMs only")
+
+    def params(suffix):
+        return {
+            "w_ih": _mlx_param(rnn, f"weight_ih_l0{suffix}", getattr(rnn, f"weight_ih_l0{suffix}"), dtype),
+            "w_hh": _mlx_param(rnn, f"weight_hh_l0{suffix}", getattr(rnn, f"weight_hh_l0{suffix}"), dtype),
+            "b_ih": _mlx_param(rnn, f"bias_ih_l0{suffix}", getattr(rnn, f"bias_ih_l0{suffix}"), dtype),
+            "b_hh": _mlx_param(rnn, f"bias_hh_l0{suffix}", getattr(rnn, f"bias_hh_l0{suffix}"), dtype),
+        }
+
+    def run(p, reverse=False):
+        steps = range(x.shape[0] - 1, -1, -1) if reverse else range(x.shape[0])
+        h = mx.zeros((x.shape[1], rnn.hidden_size), dtype=x.dtype)
+        c = mx.zeros_like(h)
+        outs = []
+        for t in steps:
+            gates = _linear(x[t], p["w_ih"], p["b_ih"]) + _linear(h, p["w_hh"], p["b_hh"])
+            i, f, g, o = mx.split(gates, 4, axis=-1)
+            i, f, o = mx.sigmoid(i), mx.sigmoid(f), mx.sigmoid(o)
+            c = f * c + i * mx.tanh(g)
+            h = o * mx.tanh(c)
+            outs.append(h)
+        if reverse:
+            outs.reverse()
+        return mx.stack(outs, axis=0)
+
+    forward = run(params(""))
+    if not rnn.bidirectional:
+        return forward
+    return mx.concatenate((forward, run(params("_reverse"), reverse=True)), axis=-1)
+
+
+def _new_lstm_module(module, x, dtype):
+    import mlx.core as mx
+
+    batch, _, nbins, nframes = x.shape
+    x = _module_forward(module.conv, x, dtype)[:, 0].transpose(2, 0, 1)
+    hidden = _lstm_forward(module.lstm, x, dtype)
+    hidden = hidden.reshape(-1, hidden.shape[-1])
+    hidden = _module_forward(module.dense, hidden, dtype)
+    return hidden.reshape(nframes, batch, 1, nbins).transpose(1, 2, 3, 0)
+
+
+def _new_base_net(module, x, dtype):
+    import mlx.core as mx
+
+    enc1 = _module_forward(module.enc1, x, dtype)
+    enc2 = _new_encoder(module.enc2, enc1, dtype)
+    enc3 = _new_encoder(module.enc3, enc2, dtype)
+    enc4 = _new_encoder(module.enc4, enc3, dtype)
+    enc5 = _new_encoder(module.enc5, enc4, dtype)
+
+    x = _new_aspp(module.aspp, enc5, dtype)
+    x = _new_decoder(module.dec4, x, enc4, dtype)
+    x = _new_decoder(module.dec3, x, enc3, dtype)
+    x = _new_decoder(module.dec2, x, enc2, dtype)
+    x = mx.concatenate((x, _new_lstm_module(module.lstm_dec2, x, dtype)), axis=1)
+    return _new_decoder(module.dec1, x, enc1, dtype)
+
+
+def _new_cascaded_net(module, x, dtype):
+    import mlx.core as mx
+
+    x = x[:, :, : module.max_bin]
+    bandwidth = x.shape[2] // 2
+    low_in = x[:, :, :bandwidth]
+    high_in = x[:, :, bandwidth:]
+
+    low1 = _module_forward(module.stg1_low_band_net, low_in, dtype)
+    high1 = _module_forward(module.stg1_high_band_net, high_in, dtype)
+    aux1 = mx.concatenate((low1, high1), axis=2)
+
+    low2 = _module_forward(module.stg2_low_band_net, mx.concatenate((low_in, low1), axis=1), dtype)
+    high2 = _module_forward(module.stg2_high_band_net, mx.concatenate((high_in, high1), axis=1), dtype)
+    aux2 = mx.concatenate((low2, high2), axis=2)
+
+    full = _module_forward(module.stg3_full_band_net, mx.concatenate((x, aux1, aux2), axis=1), dtype)
+    mask = mx.sigmoid(_conv2d_nchw(module.out, full, dtype))
+    return _replicate_pad_freq_bottom(mask, module.output_bin - mask.shape[2])
+
+
+def mlx_predict_mask_vr_mx(module, x, dtype=torch.float16):
+    if dtype not in (torch.float16, torch.float32):
+        raise TypeError("MLX full VR supports torch.float16 or torch.float32 compute dtype")
+    mx_dtype = _mlx_dtype(dtype)
+    x = x.astype(mx_dtype)
+    if isinstance(module, nets.CascadedASPPNet):
+        mask = _old_cascaded_aspp_net(module, x, dtype)
+    elif isinstance(module, nets_new.CascadedNet):
+        mask = _new_cascaded_net(module, x, dtype)
+    else:
+        raise TypeError(f"unsupported VR model for MLX full backend: {type(module).__name__}")
+    if module.offset > 0:
+        mask = mask[:, :, :, module.offset : -module.offset]
+        if mask.shape[3] <= 0:
+            raise ValueError("Window size error: h1_shape[3] must be greater than h2_shape[3]")
+    return mask

tools/pymss/modules/vocal_remover/vr_models.py

@@ -0,0 +1,116 @@
+import os
+
+
+VR_MODEL_METADATA = {
+    "10_SP-UVR-2B-32000-1.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "2band_32000"},
+    "11_SP-UVR-2B-32000-2.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "2band_32000"},
+    "12_SP-UVR-3B-44100.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "3band_44100"},
+    "13_SP-UVR-4B-44100-1.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "4band_44100"},
+    "14_SP-UVR-4B-44100-2.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "4band_44100"},
+    "15_SP-UVR-MID-44100-1.pth": {
+        "primary_stem": "Instrumental",
+        "secondary_stem": "Vocals",
+        "vr_model_param": "3band_44100_mid",
+    },
+    "16_SP-UVR-MID-44100-2.pth": {
+        "primary_stem": "Instrumental",
+        "secondary_stem": "Vocals",
+        "vr_model_param": "3band_44100_mid",
+    },
+    "17_HP-Wind_Inst-UVR.pth": {"primary_stem": "No Woodwinds", "secondary_stem": "Woodwinds", "vr_model_param": "4band_v3"},
+    "1_HP-UVR.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "4band_44100"},
+    "2_HP-UVR.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "4band_v2"},
+    "3_HP-Vocal-UVR.pth": {"primary_stem": "Vocals", "secondary_stem": "Instrumental", "vr_model_param": "4band_44100"},
+    "4_HP-Vocal-UVR.pth": {"primary_stem": "Vocals", "secondary_stem": "Instrumental", "vr_model_param": "4band_44100"},
+    "5_HP-Karaoke-UVR.pth": {
+        "primary_stem": "Instrumental",
+        "secondary_stem": "Vocals",
+        "vr_model_param": "4band_v2_sn",
+        "is_karaoke": True,
+    },
+    "6_HP-Karaoke-UVR.pth": {
+        "primary_stem": "Instrumental",
+        "secondary_stem": "Vocals",
+        "vr_model_param": "3band_44100_msb2",
+        "is_karaoke": True,
+    },
+    "7_HP2-UVR.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "3band_44100_msb2"},
+    "8_HP2-UVR.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "4band_44100"},
+    "9_HP2-UVR.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "4band_44100"},
+    "Harmonic_Noise_Separation_yxlllc.pth": {
+        "primary_stem": "No Aspiration",
+        "secondary_stem": "Aspiration",
+        "vr_model_param": "1band_sr44100_hl1024",
+    },
+    "MGM_HIGHEND_v4.pth": {
+        "primary_stem": "Instrumental",
+        "secondary_stem": "Vocals",
+        "vr_model_param": "1band_sr44100_hl1024",
+    },
+    "MGM_LOWEND_A_v4.pth": {
+        "primary_stem": "Instrumental",
+        "secondary_stem": "Vocals",
+        "vr_model_param": "1band_sr32000_hl512",
+    },
+    "MGM_LOWEND_B_v4.pth": {
+        "primary_stem": "Instrumental",
+        "secondary_stem": "Vocals",
+        "vr_model_param": "1band_sr33075_hl384",
+    },
+    "MGM_MAIN_v4.pth": {"primary_stem": "Instrumental", "secondary_stem": "Vocals", "vr_model_param": "1band_sr44100_hl512"},
+    "UVR-BVE-4B_SN-44100-1.pth": {
+        "primary_stem": "Vocals",
+        "secondary_stem": "Instrumental",
+        "vr_model_param": "4band_v3_sn",
+        "is_bv_model": True,
+        "is_bv_model_rebalanced": 0.9,
+        "nout": 64,
+        "nout_lstm": 128,
+    },
+    "UVR-De-Echo-Aggressive.pth": {
+        "primary_stem": "No Echo",
+        "secondary_stem": "Echo",
+        "vr_model_param": "4band_v3",
+        "nout": 48,
+        "nout_lstm": 128,
+    },
+    "UVR-De-Echo-Normal.pth": {
+        "primary_stem": "No Echo",
+        "secondary_stem": "Echo",
+        "vr_model_param": "4band_v3",
+        "nout": 48,
+        "nout_lstm": 128,
+    },
+    "UVR-DeEcho-DeReverb.pth": {"primary_stem": "No Reverb", "secondary_stem": "Reverb", "vr_model_param": "4band_v3"},
+    "UVR-DeNoise-Lite.pth": {
+        "primary_stem": "Noise",
+        "secondary_stem": "No Noise",
+        "vr_model_param": "1band_sr44100_hl1024",
+        "nout": 16,
+        "nout_lstm": 128,
+    },
+    "UVR-DeNoise.pth": {
+        "primary_stem": "Noise",
+        "secondary_stem": "No Noise",
+        "vr_model_param": "4band_v3",
+        "nout": 48,
+        "nout_lstm": 128,
+    },
+    "UVR-DeReverb-aufr33-jarredou_4band_v4_ms_fullband.pth": {
+        "primary_stem": "Dry",
+        "secondary_stem": "Reverb",
+        "vr_model_param": "4band_v4_ms_fullband",
+        "nout": 32,
+        "nout_lstm": 128,
+    },
+}
+
+
+def get_vr_model_metadata(model_path):
+    model_name = os.path.basename(model_path)
+    if model_name not in VR_MODEL_METADATA:
+        raise ValueError(f"Unsupported VR model: {model_name}. Only the supported UVR/VR series weights are available.")
+    data = dict(VR_MODEL_METADATA[model_name])
+    data["model_name"] = model_name
+    data["model_class"] = "VR_Models"
+    return data

tools/pymss/modules/vocal_remover/vr_separator.py

@@ -0,0 +1,459 @@
+import math
+import os
+from importlib.resources import files
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.utils.fusion import fuse_conv_bn_eval
+from tqdm import tqdm
+
+from pymss_core.modules.vocal_remover import ModelParameters, determine_model_capacity
+from pymss_core.modules.vocal_remover.uvr_lib_v5.vr_network import nets_new
+
+from .common_separator import CommonSeparator
+from .uvr_lib_v5 import spec_utils
+
+
+class _ResourceDir:
+    def __init__(self, package):
+        self._root = files(package)
+
+    def __truediv__(self, name):
+        return self._root / name
+
+    def exists(self):
+        return self._root.is_dir()
+
+    def is_dir(self):
+        return self._root.is_dir()
+
+    def iterdir(self):
+        return self._root.iterdir()
+
+    def __str__(self):
+        return str(self._root)
+
+    def __repr__(self):
+        return repr(self._root)
+
+
+VR_PARAMS_DIR = _ResourceDir("pymss_core.resources.vr_modelparams")
+
+
+def _fuse_sequential_conv_bn(module):
+    fused = 0
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Sequential):
+            new_children = []
+            child_items = list(child._modules.items())
+            i = 0
+            while i < len(child_items):
+                child_name, current = child_items[i]
+                if (
+                    i + 1 < len(child_items)
+                    and isinstance(current, nn.Conv2d)
+                    and isinstance(child_items[i + 1][1], nn.BatchNorm2d)
+                ):
+                    try:
+                        new_children.append((child_name, fuse_conv_bn_eval(current, child_items[i + 1][1])))
+                        fused += 1
+                        i += 2
+                        continue
+                    except Exception:
+                        pass
+                child_fused = _fuse_sequential_conv_bn(current)
+                fused += child_fused
+                new_children.append((child_name, current))
+                i += 1
+
+            child._modules.clear()
+            for child_name, current in new_children:
+                child.add_module(child_name, current)
+        else:
+            fused += _fuse_sequential_conv_bn(child)
+    return fused
+
+
+class VRSeparator(CommonSeparator):
+    def __init__(self, common_config, arch_config):
+        super().__init__(common_config)
+        self.model_capacity = (32, 128)
+        self.is_vr_51_model = False
+        if "nout" in self.model_data and "nout_lstm" in self.model_data:
+            self.model_capacity = (self.model_data["nout"], self.model_data["nout_lstm"])
+            self.is_vr_51_model = True
+
+        params_path = VR_PARAMS_DIR / f"{self.model_data['vr_model_param']}.json"
+        if not params_path.is_file():
+            raise FileNotFoundError(f"VR model parameter file not found: {params_path}")
+        self.model_params = ModelParameters(str(params_path))
+
+        self.enable_tta = bool(arch_config.get("enable_tta", False))
+        self.enable_post_process = bool(arch_config.get("enable_post_process", False))
+        self.post_process_threshold = float(arch_config.get("post_process_threshold", 0.2))
+        self.batch_size = int(arch_config.get("batch_size", 2))
+        self.window_size = int(arch_config.get("window_size", 512))
+        self.high_end_process = bool(arch_config.get("high_end_process", False))
+        self.use_amp = bool(arch_config.get("use_amp", True))
+        device_type = torch.device(self.torch_device).type
+        self.fuse_conv_bn = bool(arch_config.get("fuse_conv_bn", False))
+        self.use_channels_last = bool(arch_config.get("use_channels_last", False)) and device_type == "cuda"
+        self.input_high_end_h = None
+        self.input_high_end = None
+        self.aggression = float(int(arch_config.get("aggression", 5)) / 100)
+        self.aggressiveness = {
+            "value": self.aggression,
+            "split_bin": self.model_params.param["band"][1]["crop_stop"],
+            "aggr_correction": self.model_params.param.get("aggr_correction"),
+        }
+        self.model_samplerate = self.model_params.param["sr"]
+        self.model_run = None
+        self.mps_model_backend = str(arch_config.get("mps_model_backend", "torch")).lower()
+        self.mps_model_compute_dtype = self._parse_mps_model_compute_dtype(
+            arch_config.get("mps_model_compute_dtype", torch.float16)
+        )
+        if self.mps_model_backend not in ("torch", "mlx_full"):
+            raise ValueError("mps_model_backend must be 'torch' or 'mlx_full'")
+
+    @staticmethod
+    def _parse_mps_model_compute_dtype(compute_dtype):
+        if isinstance(compute_dtype, str):
+            compute_dtype = {
+                "float16": torch.float16,
+                "fp16": torch.float16,
+                "float32": torch.float32,
+                "fp32": torch.float32,
+            }.get(compute_dtype.lower(), compute_dtype)
+        if compute_dtype not in (torch.float16, torch.float32):
+            raise ValueError("mps_model_compute_dtype must be 'float16' or 'float32'")
+        return compute_dtype
+
+    def set_mps_model_backend(self, backend=None, compute_dtype=None):
+        backend = (backend or "torch").lower()
+        if backend not in ("torch", "mlx_full"):
+            raise ValueError("mps_model_backend must be 'torch' or 'mlx_full'")
+        self.mps_model_backend = backend
+        if compute_dtype is not None:
+            self.mps_model_compute_dtype = self._parse_mps_model_compute_dtype(compute_dtype)
+
+    def _use_mlx_full_forward(self, device):
+        return (
+            self.mps_model_backend == "mlx_full"
+            and torch.device(device).type == "mps"
+            and self.model_run is not None
+            and not self.model_run.training
+        )
+
+    def _store_torch_model_on_cpu_for_mlx(self):
+        return self.mps_model_backend == "mlx_full" and torch.device(self.torch_device).type == "mps"
+
+    def _predict_mask_mlx(self, x_batch_cpu):
+        import mlx.core as mx
+
+        from .vr_mlx import mlx_predict_mask_vr_mx
+
+        x_mx = mx.array(x_batch_cpu.to(dtype=self.mps_model_compute_dtype).numpy())
+        return mlx_predict_mask_vr_mx(self.model_run, x_mx, self.mps_model_compute_dtype)
+
+    def load_model(self):
+        nn_arch_sizes = [31191, 33966, 56817, 123821, 123812, 129605, 218409, 537238, 537227]
+        vr_5_1_models = [56817, 218409]
+        model_size = math.ceil(os.stat(self.model_path).st_size / 1024)
+        nn_arch_size = min(nn_arch_sizes, key=lambda size: abs(size - model_size))
+        self.logger.debug(f"VR model size: {model_size}, architecture size: {nn_arch_size}")
+
+        if nn_arch_size in vr_5_1_models or self.is_vr_51_model:
+            self.model_run = nets_new.CascadedNet(
+                self.model_params.param["bins"] * 2,
+                nn_arch_size,
+                nout=self.model_capacity[0],
+                nout_lstm=self.model_capacity[1],
+            )
+            self.is_vr_51_model = True
+        else:
+            self.model_run = determine_model_capacity(self.model_params.param["bins"] * 2, nn_arch_size)
+
+        try:
+            state_dict = torch.load(self.model_path, map_location="cpu", weights_only=True)
+        except TypeError:
+            state_dict = torch.load(self.model_path, map_location="cpu")
+        except Exception:
+            state_dict = torch.load(self.model_path, map_location="cpu", weights_only=False)
+        self.model_run.load_state_dict(state_dict)
+        self.model_run.eval()
+        if self.fuse_conv_bn:
+            fused = _fuse_sequential_conv_bn(self.model_run)
+            self.logger.debug(f"Fused {fused} VR Conv2d+BatchNorm2d pairs")
+        target_device = "cpu" if self._store_torch_model_on_cpu_for_mlx() else self.torch_device
+        self.model_run.to(target_device)
+        if self.use_channels_last:
+            self.model_run.to(memory_format=torch.channels_last)
+        self.model_run.eval()
+
+    def to(self, device):
+        self.torch_device = device
+        if self.model_run is not None:
+            self.model_run.to("cpu" if self._store_torch_model_on_cpu_for_mlx() else device)
+            if self.use_channels_last:
+                self.model_run.to(memory_format=torch.channels_last)
+        return self
+
+    def eval(self):
+        if self.model_run is not None:
+            self.model_run.eval()
+        return self
+
+    def separate_array(self, mix, sample_rate):
+        if self.model_run is None:
+            self.load_model()
+        self.primary_source = None
+        self.secondary_source = None
+        x_spec = self.loading_mix(mix, sample_rate)
+        y_spec, v_spec = self.inference_vr(x_spec, self.torch_device, self.aggressiveness)
+        y_spec = np.nan_to_num(y_spec, nan=0.0, posinf=0.0, neginf=0.0, copy=False)
+        v_spec = np.nan_to_num(v_spec, nan=0.0, posinf=0.0, neginf=0.0, copy=False)
+
+        results = {
+            self.primary_stem_name: self.process_stem(self.primary_source, y_spec),
+            self.secondary_stem_name: self.process_stem(self.secondary_source, v_spec),
+        }
+        if "Aspiration" in results:
+            aspiration = results["Aspiration"]
+            results["No Aspiration"] = aspiration[:, 1] - aspiration[:, 0]
+            results["Aspiration"] = aspiration[:, 0]
+        return results
+
+    def process_stem(self, stem_source, spec):
+        if not isinstance(stem_source, np.ndarray):
+            stem_source = self.spec_to_wav(spec).T
+            if self.model_samplerate != 44100:
+                stem_source = spec_utils.resample_audio(stem_source.T, orig_sr=self.model_samplerate, target_sr=44100).T
+        return stem_source.astype(np.float32, copy=False)
+
+    def loading_mix(self, mix, sample_rate):
+        x_wave, x_spec_s = {}, {}
+        bands_n = len(self.model_params.param["band"])
+        base_wave = self._ensure_stereo(mix)
+
+        iterator = tqdm(range(bands_n, 0, -1), leave=False, desc="Processing VR bands") if self.debug else range(bands_n, 0, -1)
+        for d in iterator:
+            bp = self.model_params.param["band"][d]
+            wav_resolution = "polyphase" if self.torch_device_mps is not None else bp["res_type"]
+            if d == bands_n:
+                x_wave[d] = self._resample_wave(base_wave, sample_rate, bp["sr"], wav_resolution)
+                x_spec_s[d] = spec_utils.wave_to_spectrogram(
+                    x_wave[d],
+                    bp["hl"],
+                    bp["n_fft"],
+                    self.model_params,
+                    band=d,
+                    is_v51_model=self.is_vr_51_model,
+                    torch_device=self.torch_device,
+                )
+            else:
+                x_wave[d] = spec_utils.resample_audio(
+                    x_wave[d + 1],
+                    orig_sr=self.model_params.param["band"][d + 1]["sr"],
+                    target_sr=bp["sr"],
+                    res_type=wav_resolution,
+                )
+                x_spec_s[d] = spec_utils.wave_to_spectrogram(
+                    x_wave[d],
+                    bp["hl"],
+                    bp["n_fft"],
+                    self.model_params,
+                    band=d,
+                    is_v51_model=self.is_vr_51_model,
+                    torch_device=self.torch_device,
+                )
+
+            if d == bands_n and self.high_end_process:
+                self.input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (
+                    self.model_params.param["pre_filter_stop"] - self.model_params.param["pre_filter_start"]
+                )
+                self.input_high_end = x_spec_s[d][:, bp["n_fft"] // 2 - self.input_high_end_h : bp["n_fft"] // 2, :]
+
+        return spec_utils.combine_spectrograms(x_spec_s, self.model_params, is_v51_model=self.is_vr_51_model)
+
+    def _ensure_stereo(self, mix):
+        mix = np.asarray(mix, dtype=np.float32)
+        if mix.ndim == 1:
+            return np.asfortranarray([mix, mix])
+        if mix.shape[0] == 2:
+            return np.asfortranarray(mix)
+        if mix.shape[-1] == 2:
+            return np.asfortranarray(mix.T)
+        return np.asfortranarray([mono := np.mean(mix, axis=0), mono])
+
+    @staticmethod
+    def _resample_wave(wave, orig_sr, target_sr, res_type):
+        return (
+            np.asfortranarray(wave)
+            if int(orig_sr) == int(target_sr)
+            else spec_utils.resample_audio(wave, orig_sr=orig_sr, target_sr=target_sr, res_type=res_type)
+        )
+
+    def inference_vr(self, x_spec, device, aggressiveness):
+        def execute(x_mag_pad, roi_size):
+            patches = (x_mag_pad.shape[2] - 2 * self.model_run.offset) // roi_size
+            x_dataset = [x_mag_pad[:, :, i * roi_size : i * roi_size + self.window_size] for i in range(patches)]
+            if not x_dataset:
+                raise ValueError("Window size error: no VR patches generated")
+
+            x_dataset = np.asarray(x_dataset)
+            mask = None
+            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
+
+                mask_batches = []
+                for i in process_batches:
+                    batch_count = min(self.batch_size, patches - i)
+                    pred = self._predict_mask_mlx(torch.from_numpy(x_dataset[i : i + batch_count]))
+                    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.progress_callback:
+                        self.progress_callback(i + batch_count, patches, "Processing VR batches")
+                return mx.concatenate(mask_batches, axis=2)[:, :, :write_pos]
+
+            with torch.inference_mode():
+                for i in process_batches:
+                    batch_count = min(self.batch_size, patches - i)
+                    x_batch_cpu = torch.from_numpy(x_dataset[i : i + batch_count])
+                    x_batch = (
+                        x_batch_cpu.to(device=device, non_blocking=True, memory_format=torch.channels_last)
+                        if self.use_channels_last
+                        else x_batch_cpu.to(device=device, non_blocking=True)
+                    )
+                    device_type = torch.device(device).type
+                    use_amp = self.use_amp and device_type in ("cuda", "mps")
+                    with torch.amp.autocast(device_type, dtype=torch.float16, enabled=use_amp):
+                        pred = self.model_run.predict_mask(x_batch)
+                    if not pred.size()[3] > 0:
+                        raise ValueError("Window size error: h1_shape[3] must be greater than h2_shape[3]")
+                    pred = pred.detach().float().permute(1, 2, 0, 3).reshape(pred.size(1), pred.size(2), -1)
+                    if mask is None:
+                        mask = torch.empty(
+                            (pred.size(0), pred.size(1), patches * pred.size(2)),
+                            dtype=pred.dtype,
+                            device=pred.device,
+                        )
+                    mask[:, :, write_pos : write_pos + pred.size(2)] = pred
+                    write_pos += pred.size(2)
+                    if self.progress_callback:
+                        self.progress_callback(i + batch_count, patches, "Processing VR batches")
+            return mask[:, :, :write_pos]
+
+        def adjust_aggr_torch(mask, is_non_accom_stem):
+            aggr = aggressiveness["value"] * 2
+            if aggr == 0:
+                return mask
+            mask = mask.clone()
+            if is_non_accom_stem:
+                aggr = 1 - aggr
+            if aggr > 10 or aggr < -10:
+                print(f"Warning: Extreme aggressiveness values detected: {aggr}")
+
+            aggr = torch.tensor([aggr, aggr], dtype=mask.dtype, device=mask.device)
+            if (correction := aggressiveness["aggr_correction"]) is not None:
+                aggr[0] += correction["left"]
+                aggr[1] += correction["right"]
+
+            split_bin = aggressiveness["split_bin"]
+            mask[:, :split_bin] = torch.pow(mask[:, :split_bin], 1 + aggr[:, None, None] / 3)
+            mask[:, split_bin:] = torch.pow(mask[:, split_bin:], 1 + aggr[:, None, None])
+            return mask
+
+        def adjust_aggr_mlx(mask, is_non_accom_stem):
+            import mlx.core as mx
+
+            aggr = aggressiveness["value"] * 2
+            if aggr == 0:
+                return mask
+            if is_non_accom_stem:
+                aggr = 1 - aggr
+            if aggr > 10 or aggr < -10:
+                print(f"Warning: Extreme aggressiveness values detected: {aggr}")
+
+            aggr = mx.array([aggr, aggr], dtype=mask.dtype)
+            if (correction := aggressiveness["aggr_correction"]) is not None:
+                correction_arr = mx.array([correction["left"], correction["right"]], dtype=mask.dtype)
+                aggr = aggr + correction_arr
+
+            split_bin = aggressiveness["split_bin"]
+            low = mx.power(mask[:, :split_bin], 1 + aggr[:, None, None] / 3)
+            high = mx.power(mask[:, split_bin:], 1 + aggr[:, None, None])
+            return mx.concatenate((low, high), axis=1)
+
+        def postprocess(mask, x_spec):
+            is_non_accom_stem = self.primary_stem_name in CommonSeparator.NON_ACCOM_STEMS
+            if self.enable_post_process:
+                if not isinstance(mask, torch.Tensor):
+                    mask = np.array(mask, copy=False)
+                else:
+                    mask = mask.cpu().numpy()
+                mask = spec_utils.adjust_aggr(mask, is_non_accom_stem, aggressiveness)
+                mask = spec_utils.merge_artifacts(mask, thres=self.post_process_threshold)
+                y_spec = mask * x_spec
+                v_spec = (1 - mask) * x_spec
+                return y_spec, v_spec
+
+            if not isinstance(mask, torch.Tensor):
+                import mlx.core as mx
+
+                mask = adjust_aggr_mlx(mask, is_non_accom_stem)
+                x_spec_mx = mx.array(x_spec)
+                y_spec = mask * x_spec_mx
+                v_spec = (1 - mask) * x_spec_mx
+                return np.array(y_spec, copy=False), np.array(v_spec, copy=False)
+
+            mask = adjust_aggr_torch(mask, is_non_accom_stem)
+            x_spec_t = torch.from_numpy(x_spec).to(device)
+            y_spec = (mask * x_spec_t).cpu().numpy()
+            v_spec = ((1 - mask) * x_spec_t).cpu().numpy()
+            return y_spec, v_spec
+
+        x_mag = np.abs(x_spec)
+        n_frame = x_mag.shape[2]
+        pad_l, pad_r, roi_size = spec_utils.make_padding(n_frame, self.window_size, self.model_run.offset)
+        x_mag_pad = np.pad(x_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")
+        max_value = x_mag_pad.max()
+        if max_value > 0:
+            x_mag_pad /= max_value
+        mask = execute(x_mag_pad, roi_size)
+
+        if self.enable_tta:
+            pad_l += roi_size // 2
+            pad_r += roi_size // 2
+            x_mag_pad = np.pad(x_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode="constant")
+            max_value = x_mag_pad.max()
+            if max_value > 0:
+                x_mag_pad /= max_value
+            mask_tta = execute(x_mag_pad, roi_size)[:, :, roi_size // 2 :]
+            mask = (mask[:, :, :n_frame] + mask_tta[:, :, :n_frame]) * 0.5
+        else:
+            mask = mask[:, :, :n_frame]
+
+        return postprocess(mask, x_spec)
+
+    def spec_to_wav(self, spec):
+        if self.high_end_process and isinstance(self.input_high_end, np.ndarray) and self.input_high_end_h:
+            input_high_end = spec_utils.mirroring("mirroring", spec, self.input_high_end, self.model_params)
+            return spec_utils.cmb_spectrogram_to_wave(
+                spec,
+                self.model_params,
+                self.input_high_end_h,
+                input_high_end,
+                is_v51_model=self.is_vr_51_model,
+                torch_device=self.torch_device,
+            )
+        return spec_utils.cmb_spectrogram_to_wave(
+            spec, self.model_params, is_v51_model=self.is_vr_51_model, torch_device=self.torch_device
+        )

tools/pymss/progress.py

@@ -0,0 +1,179 @@
+import math
+from time import time
+
+from tqdm.auto import tqdm
+
+
+def _format_progress_time(value):
+    """Format progress seconds as mm:ss or hh:mm:ss."""
+    seconds = max(0, int(round(value or 0)))
+    hours, remainder = divmod(seconds, 3600)
+    minutes, seconds = divmod(remainder, 60)
+    if hours:
+        return f"{hours:d}:{minutes:02d}:{seconds:02d}"
+    return f"{minutes:d}:{seconds:02d}"
+
+
+def _is_time_progress_message(message):
+    """Return whether a progress message represents audio seconds."""
+    return "audio" in str(message).lower()
+
+
+def _format_progress_rtf(bar):
+    """Return real-time factor text for a tqdm audio progress bar."""
+    done = float(bar.n or 0)
+    if done <= 0:
+        return "--"
+    elapsed = max(0.0, time() - getattr(bar, "start_t", time()))
+    return f"{elapsed / done:.2f}"
+
+
+def _update_time_progress_bar(bar):
+    """Update a tqdm bar with formatted elapsed/total audio time."""
+    if bar is None:
+        return
+    rtf = _format_progress_rtf(bar)
+    bar._pymss_audio_progress = f"{_format_progress_time(bar.n)}/{_format_progress_time(bar.total)}, RTF={rtf}"
+    bar.refresh()
+
+
+class _TimeProgressTqdm(tqdm):
+    """tqdm subclass with a custom audio progress field."""
+
+    @property
+    def format_dict(self):
+        data = super().format_dict
+        data["audio_progress"] = getattr(self, "_pymss_audio_progress", "")
+        return data
+
+
+class _ProgressContext:
+    """Small progress adapter used by demixing helpers."""
+
+    def __init__(
+        self,
+        pbar=False,
+        total=1,
+        callback=None,
+        done=0,
+        message="Processing audio",
+        sample_rate=None,
+    ):
+        """Initialize the progress adapter.
+
+        Args:
+            pbar (Any, optional): Whether to show a tqdm progress bar.
+                Defaults to False.
+            total (Any, optional): Total progress units. Defaults to 1.
+            callback (Any, optional): Optional callback receiving
+                ``(done, total, message)``. Defaults to None.
+            done (Any, optional): Initial completed units. Defaults to 0.
+            message (str, optional): Progress message. Defaults to
+                ``"Processing audio"``.
+            sample_rate (int | None, optional): Sample rate used to expose
+                progress in seconds. When omitted, progress values are used as
+                already provided.
+        """
+        self.enabled = bool(pbar or callback)
+        self.bar = None
+        self.callback = callback
+        self.done = done
+        self.total = total
+        self.message = message
+        self.sample_rate = int(sample_rate or 0)
+        if not self.enabled:
+            return
+        self.total = int(self.total or 1)
+        self.done = min(max(0, int(self.done or 0)), self.total)
+        if pbar:
+            bar_kwargs = {"total": self._display_total(), "desc": message, "leave": False}
+            if self.sample_rate > 0:
+                bar_kwargs.update({"unit": "", "bar_format": "{l_bar}{bar}| {audio_progress}"})
+                self.bar = _TimeProgressTqdm(**bar_kwargs)
+            else:
+                self.bar = tqdm(**bar_kwargs)
+            if self.sample_rate > 0:
+                _update_time_progress_bar(self.bar)
+        if self.bar is not None and self.done:
+            self.bar.update(self._display_value(self.done))
+            if self.sample_rate > 0:
+                _update_time_progress_bar(self.bar)
+        self.emit()
+
+    def _display_value(self, value):
+        """Return progress value exposed to callbacks and progress bars."""
+        if self.sample_rate <= 0:
+            return int(value)
+        if int(value) >= self.total:
+            return self._display_total()
+        return min(self._display_total(), int(int(value) // self.sample_rate))
+
+    def _display_total(self):
+        """Return total seconds for the current progress unit."""
+        if self.sample_rate <= 0:
+            return self.total
+        return max(1, int(math.ceil(self.total / self.sample_rate)))
+
+    def emit(self, done=None):
+        """Emit a progress update."""
+        if not self.enabled:
+            return
+        if done is not None:
+            next_done = min(max(0, int(done)), self.total)
+            if self.bar is not None:
+                self.bar.update(self._display_value(next_done) - self._display_value(self.done))
+                if self.sample_rate > 0:
+                    _update_time_progress_bar(self.bar)
+            self.done = next_done
+        if self.callback is None:
+            return
+        self.callback(self._display_value(self.done), self._display_total(), self.message)
+
+    def update(self, amount):
+        """Advance progress by ``amount`` internal units."""
+        if not self.enabled:
+            return
+        amount = int(amount)
+        self.emit(self.done + amount)
+
+    def close(self):
+        """Close the progress bar when present."""
+        if not self.enabled:
+            return
+        if self.bar:
+            self.bar.close()
+
+
+class _CliInferenceProgress:
+    """CLI callback adapter for inference progress updates."""
+
+    def __init__(self):
+        self._bar = None
+        self._message = None
+        self._total = None
+
+    def __call__(self, done, total, message):
+        total = max(1, int(total or 1))
+        done = max(0, min(int(done), total))
+        if self._bar is None or self._message != message or self._total != total or done < self._bar.n:
+            self.close()
+            self._message = message
+            self._total = total
+            bar_kwargs = {"total": total, "desc": message, "leave": False, "mininterval": 0, "miniters": 1}
+            if _is_time_progress_message(message):
+                bar_kwargs.update({"unit": "", "bar_format": "{l_bar}{bar}| {audio_progress}"})
+                self._bar = _TimeProgressTqdm(**bar_kwargs)
+            else:
+                self._bar = tqdm(**bar_kwargs)
+            if _is_time_progress_message(message):
+                _update_time_progress_bar(self._bar)
+        if done != self._bar.n:
+            self._bar.update(done - self._bar.n)
+            if _is_time_progress_message(message):
+                _update_time_progress_bar(self._bar)
+
+    def close(self):
+        """Close the active CLI progress bar."""
+        if self._bar is not None:
+            self._bar.close()
+            self._bar = None

tools/pymss/resources/vr_modelparams/1band_sr16000_hl512.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 16000,
+			"hl": 512,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 16000,
+	"pre_filter_start": 1023,
+	"pre_filter_stop": 1024
+}

tools/pymss/resources/vr_modelparams/1band_sr32000_hl512.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 32000,
+			"hl": 512,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 32000,
+	"pre_filter_start": 1000,
+	"pre_filter_stop": 1021
+}

tools/pymss/resources/vr_modelparams/1band_sr33075_hl384.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 33075,
+			"hl": 384,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 33075,
+	"pre_filter_start": 1000,
+	"pre_filter_stop": 1021
+}

tools/pymss/resources/vr_modelparams/1band_sr44100_hl1024.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 1024,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 1023,
+	"pre_filter_stop": 1024
+}

tools/pymss/resources/vr_modelparams/1band_sr44100_hl256.json

@@ -0,0 +1,19 @@
+{
+	"bins": 256,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 0,
+			"crop_stop": 256,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 256,
+	"pre_filter_stop": 256
+}

tools/pymss/resources/vr_modelparams/1band_sr44100_hl512.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 1024,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 1023,
+	"pre_filter_stop": 1024
+}

tools/pymss/resources/vr_modelparams/1band_sr44100_hl512_cut.json

@@ -0,0 +1,19 @@
+{
+	"bins": 1024,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 2048,
+			"crop_start": 0,
+			"crop_stop": 700,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 1023,
+	"pre_filter_stop": 700
+}

tools/pymss/resources/vr_modelparams/1band_sr44100_hl512_nf1024.json

@@ -0,0 +1,19 @@
+{
+	"bins": 512,
+	"unstable_bins": 0,
+	"reduction_bins": 0,
+	"band": {
+		"1": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 512,
+			"hpf_start": -1,
+			"res_type": "sinc_best"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 511,
+	"pre_filter_stop": 512
+}

tools/pymss/resources/vr_modelparams/2band_32000.json

@@ -0,0 +1,30 @@
+{
+	"bins": 768,
+	"unstable_bins": 7,
+	"reduction_bins": 705,
+	"band": {
+		"1": {
+			"sr": 6000,
+			"hl": 66,
+			"n_fft": 512,
+			"crop_start": 0,
+			"crop_stop": 240,
+			"lpf_start": 60,
+			"lpf_stop": 118,
+			"res_type": "sinc_fastest"
+		},
+		"2": {
+			"sr": 32000,
+			"hl": 352,
+			"n_fft": 1024,
+			"crop_start": 22,
+			"crop_stop": 505,
+			"hpf_start": 44,
+			"hpf_stop": 23,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 32000,
+	"pre_filter_start": 710,
+	"pre_filter_stop": 731
+}

tools/pymss/resources/vr_modelparams/2band_44100_lofi.json

@@ -0,0 +1,30 @@
+{
+	"bins": 512,
+	"unstable_bins": 7,
+	"reduction_bins": 510,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 160,
+			"n_fft": 768,
+			"crop_start": 0,
+			"crop_stop": 192,
+			"lpf_start": 41,
+			"lpf_stop": 139,
+			"res_type": "sinc_fastest"
+		},
+		"2": {
+			"sr": 44100,
+			"hl": 640,
+			"n_fft": 1024,
+			"crop_start": 10,
+			"crop_stop": 320,
+			"hpf_start": 47,
+			"hpf_stop": 15,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 510,
+	"pre_filter_stop": 512
+}

tools/pymss/resources/vr_modelparams/2band_48000.json

@@ -0,0 +1,30 @@
+{
+	"bins": 768,
+	"unstable_bins": 7,
+	"reduction_bins": 705,
+	"band": {
+		"1": {
+			"sr": 6000,
+			"hl": 66,
+			"n_fft": 512,
+			"crop_start": 0,
+			"crop_stop": 240,
+			"lpf_start": 60,
+			"lpf_stop": 240,
+			"res_type": "sinc_fastest"
+		},
+		"2": {
+			"sr": 48000,
+			"hl": 528,
+			"n_fft": 1536,
+			"crop_start": 22,
+			"crop_stop": 505,
+			"hpf_start": 82,
+			"hpf_stop": 22,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 48000,
+	"pre_filter_start": 710,
+	"pre_filter_stop": 731
+}

tools/pymss/resources/vr_modelparams/3band_44100.json

@@ -0,0 +1,42 @@
+{
+	"bins": 768,
+	"unstable_bins": 5,
+	"reduction_bins": 733,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 768,
+			"crop_start": 0,
+			"crop_stop": 278,
+			"lpf_start": 28,
+			"lpf_stop": 140,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 768,
+			"crop_start": 14,
+			"crop_stop": 322,
+			"hpf_start": 70,
+			"hpf_stop": 14,
+			"lpf_start": 283,
+			"lpf_stop": 314,
+			"res_type": "polyphase"
+		},	
+		"3": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 131,
+			"crop_stop": 313,
+			"hpf_start": 154,
+			"hpf_stop": 141,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 757,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/3band_44100_mid.json

@@ -0,0 +1,43 @@
+{
+	"mid_side": true,
+	"bins": 768,
+	"unstable_bins": 5,
+	"reduction_bins": 733,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 768,
+			"crop_start": 0,
+			"crop_stop": 278,
+			"lpf_start": 28,
+			"lpf_stop": 140,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 768,
+			"crop_start": 14,
+			"crop_stop": 322,
+			"hpf_start": 70,
+			"hpf_stop": 14,
+			"lpf_start": 283,
+			"lpf_stop": 314,
+			"res_type": "polyphase"
+		},	
+		"3": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 131,
+			"crop_stop": 313,
+			"hpf_start": 154,
+			"hpf_stop": 141,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 757,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/3band_44100_msb2.json

@@ -0,0 +1,43 @@
+{
+	"mid_side_b2": true,
+	"bins": 640,
+	"unstable_bins": 7,
+	"reduction_bins": 565,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 108,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 187,
+			"lpf_start": 92,
+			"lpf_stop": 186,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 22050,
+			"hl": 216,
+			"n_fft": 768,
+			"crop_start": 0,
+			"crop_stop": 212,
+			"hpf_start": 68,
+			"hpf_stop": 34,
+			"lpf_start": 174,
+			"lpf_stop": 209,
+			"res_type": "polyphase"
+		},	
+		"3": {
+			"sr": 44100,
+			"hl": 432,
+			"n_fft": 640,
+			"crop_start": 66,
+			"crop_stop": 307,
+			"hpf_start": 86,
+			"hpf_stop": 72,
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 639,
+	"pre_filter_stop": 640
+}

tools/pymss/resources/vr_modelparams/4band_44100.json

@@ -0,0 +1,54 @@
+{
+	"bins": 768,
+	"unstable_bins": 7,
+	"reduction_bins": 668,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 186,
+			"lpf_start": 37,
+			"lpf_stop": 73,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 512,
+			"crop_start": 4,
+			"crop_stop": 185,			
+			"hpf_start": 36,
+			"hpf_stop": 18,
+			"lpf_start": 93,
+			"lpf_stop": 185,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 46,
+			"crop_stop": 186,
+			"hpf_start": 93,
+			"hpf_stop": 46,
+			"lpf_start": 164,
+			"lpf_stop": 186,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 121,
+			"crop_stop": 382,
+			"hpf_start": 138,
+			"hpf_stop": 123,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 740,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/4band_44100_mid.json

@@ -0,0 +1,55 @@
+{
+	"bins": 768,
+	"unstable_bins": 7,
+	"mid_side": true,
+	"reduction_bins": 668,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 186,
+			"lpf_start": 37,
+			"lpf_stop": 73,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 512,
+			"crop_start": 4,
+			"crop_stop": 185,			
+			"hpf_start": 36,
+			"hpf_stop": 18,
+			"lpf_start": 93,
+			"lpf_stop": 185,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 46,
+			"crop_stop": 186,
+			"hpf_start": 93,
+			"hpf_stop": 46,
+			"lpf_start": 164,
+			"lpf_stop": 186,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 121,
+			"crop_stop": 382,
+			"hpf_start": 138,
+			"hpf_stop": 123,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 740,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/4band_44100_msb.json

@@ -0,0 +1,55 @@
+{
+	"mid_side_b": true,
+	"bins": 768,
+	"unstable_bins": 7,
+	"reduction_bins": 668,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 186,
+			"lpf_start": 37,
+			"lpf_stop": 73,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 512,
+			"crop_start": 4,
+			"crop_stop": 185,			
+			"hpf_start": 36,
+			"hpf_stop": 18,
+			"lpf_start": 93,
+			"lpf_stop": 185,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 46,
+			"crop_stop": 186,
+			"hpf_start": 93,
+			"hpf_stop": 46,
+			"lpf_start": 164,
+			"lpf_stop": 186,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 121,
+			"crop_stop": 382,
+			"hpf_start": 138,
+			"hpf_stop": 123,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 740,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/4band_44100_msb2.json

@@ -0,0 +1,55 @@
+{
+	"mid_side_b": true,
+	"bins": 768,
+	"unstable_bins": 7,
+	"reduction_bins": 668,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 186,
+			"lpf_start": 37,
+			"lpf_stop": 73,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 512,
+			"crop_start": 4,
+			"crop_stop": 185,			
+			"hpf_start": 36,
+			"hpf_stop": 18,
+			"lpf_start": 93,
+			"lpf_stop": 185,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 46,
+			"crop_stop": 186,
+			"hpf_start": 93,
+			"hpf_stop": 46,
+			"lpf_start": 164,
+			"lpf_stop": 186,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 121,
+			"crop_stop": 382,
+			"hpf_start": 138,
+			"hpf_stop": 123,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 740,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/4band_44100_reverse.json

@@ -0,0 +1,55 @@
+{
+	"reverse": true,
+	"bins": 768,
+	"unstable_bins": 7,
+	"reduction_bins": 668,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 186,
+			"lpf_start": 37,
+			"lpf_stop": 73,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 512,
+			"crop_start": 4,
+			"crop_stop": 185,			
+			"hpf_start": 36,
+			"hpf_stop": 18,
+			"lpf_start": 93,
+			"lpf_stop": 185,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 46,
+			"crop_stop": 186,
+			"hpf_start": 93,
+			"hpf_stop": 46,
+			"lpf_start": 164,
+			"lpf_stop": 186,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 121,
+			"crop_stop": 382,
+			"hpf_start": 138,
+			"hpf_stop": 123,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 740,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/4band_44100_sw.json

@@ -0,0 +1,55 @@
+{
+	"stereo_w": true,
+	"bins": 768,
+	"unstable_bins": 7,
+	"reduction_bins": 668,
+	"band": {
+		"1": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 1024,
+			"crop_start": 0,
+			"crop_stop": 186,
+			"lpf_start": 37,
+			"lpf_stop": 73,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 11025,
+			"hl": 128,
+			"n_fft": 512,
+			"crop_start": 4,
+			"crop_stop": 185,			
+			"hpf_start": 36,
+			"hpf_stop": 18,
+			"lpf_start": 93,
+			"lpf_stop": 185,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 22050,
+			"hl": 256,
+			"n_fft": 512,
+			"crop_start": 46,
+			"crop_stop": 186,
+			"hpf_start": 93,
+			"hpf_stop": 46,
+			"lpf_start": 164,
+			"lpf_stop": 186,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 512,
+			"n_fft": 768,
+			"crop_start": 121,
+			"crop_stop": 382,
+			"hpf_start": 138,
+			"hpf_stop": 123,
+			"res_type": "sinc_medium"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 740,
+	"pre_filter_stop": 768
+}

tools/pymss/resources/vr_modelparams/4band_v2.json

@@ -0,0 +1,54 @@
+{
+	"bins": 672,
+	"unstable_bins": 8,
+	"reduction_bins": 637,
+	"band": {
+		"1": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 640,
+			"crop_start": 0,
+			"crop_stop": 85,
+			"lpf_start": 25,
+			"lpf_stop": 53,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 320,
+			"crop_start": 4,
+			"crop_stop": 87,
+			"hpf_start": 25,
+			"hpf_stop": 12,
+			"lpf_start": 31,
+			"lpf_stop": 62,
+			"res_type": "polyphase"
+		},		
+		"3": {
+			"sr": 14700,
+			"hl": 160,
+			"n_fft": 512,
+			"crop_start": 17,
+			"crop_stop": 216,
+			"hpf_start": 48,
+			"hpf_stop": 24,
+			"lpf_start": 139,
+			"lpf_stop": 210,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 480,
+			"n_fft": 960,
+			"crop_start": 78,
+			"crop_stop": 383,
+			"hpf_start": 130,
+			"hpf_stop": 86,
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 668,
+	"pre_filter_stop": 672
+}

tools/pymss/resources/vr_modelparams/4band_v2_sn.json

@@ -0,0 +1,55 @@
+{
+	"bins": 672,
+	"unstable_bins": 8,
+	"reduction_bins": 637,
+	"band": {
+		"1": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 640,
+			"crop_start": 0,
+			"crop_stop": 85,
+			"lpf_start": 25,
+			"lpf_stop": 53,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 320,
+			"crop_start": 4,
+			"crop_stop": 87,
+			"hpf_start": 25,
+			"hpf_stop": 12,
+			"lpf_start": 31,
+			"lpf_stop": 62,
+			"res_type": "polyphase"
+		},		
+		"3": {
+			"sr": 14700,
+			"hl": 160,
+			"n_fft": 512,
+			"crop_start": 17,
+			"crop_stop": 216,
+			"hpf_start": 48,
+			"hpf_stop": 24,
+			"lpf_start": 139,
+			"lpf_stop": 210,
+			"res_type": "polyphase"
+		},	
+		"4": {
+			"sr": 44100,
+			"hl": 480,
+			"n_fft": 960,
+			"crop_start": 78,
+			"crop_stop": 383,
+			"hpf_start": 130,
+			"hpf_stop": 86,
+			"convert_channels": "stereo_n",
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 668,
+	"pre_filter_stop": 672
+}

tools/pymss/resources/vr_modelparams/4band_v3.json

@@ -0,0 +1,54 @@
+{
+	"bins": 672,
+	"unstable_bins": 8,
+	"reduction_bins": 530,
+	"band": {
+		"1": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 640,
+			"crop_start": 0,
+			"crop_stop": 85,
+			"lpf_start": 25,
+			"lpf_stop": 53,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 320,
+			"crop_start": 4,
+			"crop_stop": 87,
+			"hpf_start": 25,
+			"hpf_stop": 12,
+			"lpf_start": 31,
+			"lpf_stop": 62,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 14700,
+			"hl": 160,
+			"n_fft": 512,
+			"crop_start": 17,
+			"crop_stop": 216,
+			"hpf_start": 48,
+			"hpf_stop": 24,
+			"lpf_start": 139,
+			"lpf_stop": 210,
+			"res_type": "polyphase"
+		},
+		"4": {
+			"sr": 44100,
+			"hl": 480,
+			"n_fft": 960,
+			"crop_start": 78,
+			"crop_stop": 383,
+			"hpf_start": 130,
+			"hpf_stop": 86,
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 668,
+	"pre_filter_stop": 672
+}

tools/pymss/resources/vr_modelparams/4band_v3_sn.json

@@ -0,0 +1,55 @@
+{
+	"n_bins": 672,
+	"unstable_bins": 8,
+	"stable_bins": 530,
+	"band": {
+		"1": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 640,
+			"crop_start": 0,
+			"crop_stop": 85,
+			"lpf_start": 25,
+			"lpf_stop": 53,
+			"res_type": "polyphase"
+		},
+		"2": {
+			"sr": 7350,
+			"hl": 80,
+			"n_fft": 320,
+			"crop_start": 4,
+			"crop_stop": 87,
+			"hpf_start": 25,
+			"hpf_stop": 12,
+			"lpf_start": 31,
+			"lpf_stop": 62,
+			"res_type": "polyphase"
+		},
+		"3": {
+			"sr": 14700,
+			"hl": 160,
+			"n_fft": 512,
+			"crop_start": 17,
+			"crop_stop": 216,
+			"hpf_start": 48,
+			"hpf_stop": 24,
+			"lpf_start": 139,
+			"lpf_stop": 210,
+			"res_type": "polyphase"
+		},
+		"4": {
+			"sr": 44100,
+			"hl": 480,
+			"n_fft": 960,
+			"crop_start": 78,
+			"crop_stop": 383,
+			"hpf_start": 130,
+			"hpf_stop": 86,
+			"convert_channels": "stereo_n",
+			"res_type": "kaiser_fast"
+		}
+	},
+	"sr": 44100,
+	"pre_filter_start": 668,
+	"pre_filter_stop": 672
+}

tools/pymss/resources/vr_modelparams/4band_v4_ms_fullband.json

@@ -0,0 +1,58 @@
+{
+    "n_bins": 896,
+    "unstable_bins": 9,
+    "stable_bins": 530,
+    "band": {
+        "1": {
+            "sr": 7350,
+            "hl": 96,
+            "n_fft": 768,
+            "crop_start": 0,
+            "crop_stop": 102,
+            "lpf_start": 30,
+            "lpf_stop": 62,
+            "res_type": "polyphase",
+            "convert_channels": "mid_side"
+        },
+        "2": {
+            "sr": 7350,
+            "hl": 96,
+            "n_fft": 384,
+            "crop_start": 5,
+            "crop_stop": 104,
+            "hpf_start": 30,
+            "hpf_stop": 14,
+            "lpf_start": 37,
+            "lpf_stop": 73,
+            "res_type": "polyphase",
+			"convert_channels": "mid_side"
+        },
+        "3": {
+            "sr": 14700,
+            "hl": 192,
+            "n_fft": 640,
+            "crop_start": 20,
+            "crop_stop": 259,
+            "hpf_start": 58,
+            "hpf_stop": 29,
+            "lpf_start": 191,
+            "lpf_stop": 262,
+            "res_type": "polyphase",
+			"convert_channels": "mid_side"
+        },
+        "4": {
+            "sr": 44100,
+            "hl": 576,
+            "n_fft": 1152,
+            "crop_start": 119,
+            "crop_stop": 575,
+            "hpf_start": 157,
+            "hpf_stop": 110,
+            "res_type": "kaiser_fast",
+			"convert_channels": "mid_side"
+        }
+    },
+    "sr": 44100,
+    "pre_filter_start": -1,
+    "pre_filter_stop": -1
+}

tools/pymss/server/__init__.py

@@ -0,0 +1,36 @@
+from .config import ServerConfig
+
+
+def create_app(config):
+    """Create the FastAPI application.
+
+    Args:
+        config (AttrDict | dict): Loaded pymss configuration.
+
+    Returns:
+        FastAPI: Configured application instance.
+
+    Example:
+        >>> app = create_app()"""
+    from .app import create_app as _create_app
+
+    return _create_app(config)
+
+
+def run_server(config):
+    """Run the pymss HTTP server.
+
+    Args:
+        config (AttrDict | dict): Loaded pymss configuration.
+
+    Returns:
+        None: Runs until the server stops.
+
+    Example:
+        >>> run_server()"""
+    from .app import run_server as _run_server
+
+    return _run_server(config)
+
+
+__all__ = ("ServerConfig", "create_app", "run_server")

tools/pymss/server/audio.py

@@ -0,0 +1,363 @@
+from __future__ import annotations
+
+import base64
+import io
+import json
+import os
+import re
+import tempfile
+import time
+import uuid
+import zipfile
+
+import numpy as np
+
+from ..audio_io import save_audio
+from .errors import APIError
+
+
+PCM_FORMATS = {
+    "pcm_f32le": (np.dtype("<f4"), 4),
+    "pcm_s16le": (np.dtype("<i2"), 2),
+}
+OUTPUT_FORMATS = {"pcm_f32le", "wav", "flac"}
+RESPONSE_FORMATS = {"json", "zip"}
+
+
+def parse_int(value, param, code="invalid_request"):
+    """Parse an integer request parameter and raise an APIError on failure.
+
+    Args:
+        value (Any): Value value.
+        param (str | None): Param value.
+        code (str, optional): Code value. Defaults to 'invalid_request'.
+
+    Returns:
+        Any: Parsed value."""
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        raise APIError(400, code, f"Invalid integer for {param}.", param=param)
+
+
+def normalize_stems(value, valid_instruments):
+    """Normalize requested stem names against the loaded model instruments.
+
+    Args:
+        value (Any): Value value.
+        valid_instruments (Any): Valid instruments value.
+
+    Returns:
+        Any: Computed result."""
+    if value is None:
+        return None
+    if isinstance(value, str):
+        raw_stems = value.split(",")
+    elif isinstance(value, (list, tuple)):
+        raw_stems = value
+    else:
+        raise APIError(400, "invalid_stem", "stems must be a string or array.", param="stems")
+
+    lower_map = {stem.lower(): stem for stem in valid_instruments}
+    selected = []
+    seen = set()
+    for raw in raw_stems:
+        stem = str(raw).strip()
+        if not stem:
+            continue
+        canonical = lower_map.get(stem.lower())
+        if canonical is None:
+            raise APIError(400, "invalid_stem", f"Invalid stem {stem!r}. Valid stems: {list(valid_instruments)}", param="stems")
+        if canonical.lower() in seen:
+            continue
+        seen.add(canonical.lower())
+        selected.append(canonical)
+    return selected or None
+
+
+def validate_common_options(response_format, output_audio_format):
+    """Validate response and output audio format options.
+
+    Args:
+        response_format (str): Response format value.
+        output_audio_format (str): Output audio format value.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    if response_format not in RESPONSE_FORMATS:
+        raise APIError(
+            400,
+            "invalid_response_format",
+            f"Unsupported response_format {response_format!r}.",
+            param="response_format",
+        )
+    if output_audio_format not in OUTPUT_FORMATS:
+        raise APIError(
+            400,
+            "invalid_output_audio_format",
+            f"Unsupported output_audio_format {output_audio_format!r}.",
+            param="output_audio_format",
+        )
+    if response_format == "json" and output_audio_format != "pcm_f32le":
+        raise APIError(
+            400,
+            "invalid_output_audio_format",
+            "response_format='json' only supports output_audio_format='pcm_f32le'.",
+            param="output_audio_format",
+        )
+
+
+def decode_pcm(raw, audio_format, sample_rate, channels, expected_sample_rate, max_audio_seconds):
+    """Decode raw PCM request bytes into a channel-first float32 array.
+
+    Args:
+        raw (bytes): Raw value.
+        audio_format (Any): Audio format value.
+        sample_rate (int): Audio sample rate in Hz.
+        channels (int): Channels value.
+        expected_sample_rate (Any): Expected sample rate value.
+        max_audio_seconds (Any): Max audio seconds value.
+
+    Returns:
+        Any: Computed result."""
+    if audio_format not in PCM_FORMATS:
+        raise APIError(400, "invalid_audio_format", f"Unsupported audio format {audio_format!r}.", param="format")
+    if channels not in (1, 2):
+        raise APIError(400, "invalid_channel_count", "channels must be 1 or 2.", param="channels")
+    if sample_rate != expected_sample_rate:
+        raise APIError(
+            400,
+            "invalid_sample_rate",
+            f"sample_rate must be {expected_sample_rate}.",
+            param="sample_rate",
+        )
+    if not raw:
+        raise APIError(400, "empty_audio", "Decoded PCM audio is empty.", param="input.data")
+
+    dtype, sample_width = PCM_FORMATS[audio_format]
+    frame_width = channels * sample_width
+    if len(raw) % frame_width:
+        raise APIError(
+            400,
+            "invalid_audio_length",
+            "PCM bytes length is not aligned to format and channels.",
+            param="input.data",
+        )
+
+    frames = len(raw) // frame_width
+    seconds = frames / float(sample_rate)
+    if max_audio_seconds and seconds > max_audio_seconds:
+        raise APIError(413, "request_too_large", f"Audio duration exceeds {max_audio_seconds} seconds.")
+
+    data = np.frombuffer(raw, dtype=dtype)
+    if audio_format == "pcm_f32le":
+        if not np.isfinite(data).all():
+            raise APIError(400, "invalid_audio_data", "pcm_f32le audio contains NaN or Inf.", param="input.data")
+        float_data = data.astype(np.float32, copy=False)
+    else:
+        float_data = data.astype(np.float32) / 32768.0
+
+    frame_data = float_data.reshape(-1, channels)
+    mix = frame_data[:, 0] if channels == 1 else frame_data.T
+    return np.ascontiguousarray(mix), seconds
+
+
+def audio_to_interleaved_f32(audio):
+    """Convert channel-first audio to interleaved float32 samples.
+
+    Args:
+        audio (np.ndarray): Audio samples.
+
+    Returns:
+        Any: Computed result."""
+    array = np.asarray(audio, dtype=np.float32)
+    if array.ndim == 1:
+        normalized = np.ascontiguousarray(array)
+        return normalized, 1
+    if array.ndim != 2:
+        raise ValueError(f"Expected mono or stereo audio, got shape {array.shape}")
+    if array.shape[1] in (1, 2):
+        normalized = np.ascontiguousarray(array)
+        return normalized, int(array.shape[1])
+    if array.shape[0] in (1, 2):
+        normalized = np.ascontiguousarray(array.T)
+        return normalized, int(array.shape[0])
+    raise ValueError(f"Expected mono or stereo audio, got shape {array.shape}")
+
+
+def f32le_bytes(audio):
+    """Serialize audio as little-endian float32 PCM bytes.
+
+    Args:
+        audio (np.ndarray): Audio samples.
+
+    Returns:
+        Any: Computed result."""
+    array, channels = audio_to_interleaved_f32(audio)
+    return np.asarray(array, dtype="<f4").tobytes(), channels
+
+
+def _safe_slug(value):
+    """Implement the safe slug helper.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Computed result."""
+    slug = re.sub(r"[^a-z0-9_-]+", "-", str(value).lower()).strip("-._")
+    return slug or "stem"
+
+
+def _filename(index, stem, output_audio_format):
+    """Implement the filename helper.
+
+    Args:
+        index (Any): Index value.
+        stem (str): Stem value.
+        output_audio_format (str): Output audio format value.
+
+    Returns:
+        Any: Computed result."""
+    extension = "f32le" if output_audio_format == "pcm_f32le" else output_audio_format
+    return f"{index:04d}-{_safe_slug(stem)}.{extension}"
+
+
+def _encode_container(audio, sample_rate, output_format, audio_params):
+    """Encode container.
+
+    Args:
+        audio (np.ndarray): Audio samples.
+        sample_rate (int): Audio sample rate in Hz.
+        output_format (str | None): Output format such as wav, flac, mp3, or m4a.
+        audio_params (dict | None): Encoding options for the output audio format.
+
+    Returns:
+        Any: Computed result."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = os.path.join(tmpdir, f"audio.{output_format}")
+        save_audio(path, audio, sample_rate, output_format, audio_params)
+        with open(path, "rb") as f:
+            return f.read()
+
+
+def ordered_results(results, stems, instruments):
+    """Return separation results in requested stem order.
+
+    Args:
+        results (dict): Results value.
+        stems (Sequence[str] | None): Requested output stem names.
+        instruments (Sequence[str] | None): Instruments value.
+
+    Returns:
+        Any: Computed result."""
+    order = stems or list(instruments)
+    missing = [stem for stem in order if stem not in results]
+    if missing:
+        raise APIError(
+            500,
+            "separation_failed",
+            f"Separation did not return stem(s): {missing}",
+            error_type="server_error",
+        )
+    return [(stem, results[stem]) for stem in order]
+
+
+def json_response(loaded, model, results, stems, input_seconds):
+    """Build a JSON separation response with base64 audio payloads.
+
+    Args:
+        loaded (LoadedModel): Loaded value.
+        model (str): Model value.
+        results (dict): Results value.
+        stems (Sequence[str] | None): Requested output stem names.
+        input_seconds (Any): Input seconds value.
+
+    Returns:
+        Any: Computed result."""
+    outputs = []
+    for stem, audio in ordered_results(results, stems, loaded.instruments):
+        raw, channels = f32le_bytes(audio)
+        outputs.append(
+            {
+                "stem": stem,
+                "audio": {
+                    "format": "pcm_f32le",
+                    "sample_rate": loaded.sample_rate,
+                    "channels": channels,
+                    "data": base64.b64encode(raw).decode("ascii"),
+                },
+            }
+        )
+    created = int(time.time())
+    return {
+        "id": "sep_" + uuid.uuid4().hex,
+        "object": "audio.separation",
+        "created": created,
+        "model": model,
+        "outputs": outputs,
+        "metadata": {
+            "input_seconds": input_seconds,
+            "output_stems": [item["stem"] for item in outputs],
+            "device": loaded.device,
+        },
+        "usage": {
+            "type": "duration",
+            "seconds": input_seconds,
+        },
+    }
+
+
+def zip_response(loaded, model, results, stems, input_seconds, output_audio_format):
+    """Build a ZIP separation response with encoded audio files.
+
+    Args:
+        loaded (LoadedModel): Loaded value.
+        model (str): Model value.
+        results (dict): Results value.
+        stems (Sequence[str] | None): Requested output stem names.
+        input_seconds (Any): Input seconds value.
+        output_audio_format (str): Output audio format value.
+
+    Returns:
+        Any: Computed result."""
+    output_items = []
+    archive = io.BytesIO()
+    with zipfile.ZipFile(archive, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for index, (stem, audio) in enumerate(ordered_results(results, stems, loaded.instruments), start=1):
+            filename = _filename(index, stem, output_audio_format)
+            if output_audio_format == "pcm_f32le":
+                content, channels = f32le_bytes(audio)
+                format_name = "pcm_f32le"
+            else:
+                content = _encode_container(
+                    audio,
+                    loaded.sample_rate,
+                    output_audio_format,
+                    loaded.audio_params,
+                )
+                _, channels = audio_to_interleaved_f32(audio)
+                format_name = output_audio_format
+            zf.writestr(filename, content)
+            output_items.append(
+                {
+                    "stem": stem,
+                    "filename": filename,
+                    "format": format_name,
+                    "sample_rate": loaded.sample_rate,
+                    "channels": channels,
+                }
+            )
+
+        manifest = {
+            "id": "sep_" + uuid.uuid4().hex,
+            "object": "audio.separation",
+            "model": model,
+            "outputs": output_items,
+            "usage": {
+                "type": "duration",
+                "seconds": input_seconds,
+            },
+        }
+        zf.writestr("manifest.json", json.dumps(manifest, ensure_ascii=False, indent=2))
+    return archive.getvalue()

tools/pymss/server/config.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ServerConfig:
+    """Runtime configuration for the pymss HTTP server."""
+
+    model: str | None = None
+    model_dir: str | None = None
+    source: str = "modelscope"
+    endpoint: str | None = None
+    device: str = "auto"
+    device_ids: list[int] = field(default_factory=lambda: [0])
+    api_key: str | None = None
+    host: str = "127.0.0.1"
+    port: int = 8000
+    debug: bool = False
+    inference_params: dict = field(default_factory=dict)
+    max_audio_seconds: float = 600.0
+    max_request_bytes: int = 536870912
+    max_queue_size: int = 8
+    request_timeout_seconds: float = 0.0
+    webui: bool = False

tools/pymss/server/errors.py

@@ -0,0 +1,29 @@
+class APIError(Exception):
+    """Structured HTTP API error used by the pymss server.
+
+    Args:
+        status_code (int): Status code value.
+        code (str): Code value.
+        message (str): Message value.
+        param (str | None, optional): Param value. Defaults to None.
+        error_type (str, optional): Error type value. Defaults to 'invalid_request_error'.
+    """
+
+    def __init__(self, status_code, code, message, param=None, error_type="invalid_request_error"):
+        """Initialize the instance.
+
+        Args:
+            status_code (int): Status code value.
+            code (str): Code value.
+            message (str): Message value.
+            param (str | None, optional): Param value. Defaults to None.
+            error_type (str, optional): Error type value. Defaults to 'invalid_request_error'.
+
+        Returns:
+            None: This method completes for its side effects."""
+        super().__init__(message)
+        self.status_code = status_code
+        self.code = code
+        self.message = message
+        self.param = param
+        self.error_type = error_type

tools/pymss/server/models.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+from ..model_download import remote_url
+from ..model_registry import (
+    auxiliary_paths_for,
+    config_path_for,
+    get_model_entry,
+    list_models,
+    model_path_for,
+    model_root,
+)
+
+
+def _bool_query(value, *, default=False):
+    """Implement the bool query helper.
+
+    Args:
+        value (Any): Value value.
+        default (Any, optional): Default value. Defaults to False.
+
+    Returns:
+        Any: Computed result."""
+    if value is None:
+        return default
+    value = str(value).strip().lower()
+    if value in {"1", "true", "yes", "on"}:
+        return True
+    if value in {"0", "false", "no", "off"}:
+        return False
+    raise ValueError("Expected boolean value")
+
+
+def parse_supported_filter(value):
+    """Parse the model catalog supported filter.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Parsed value."""
+    if value is None:
+        return True
+    value = str(value).strip().lower()
+    if value == "all":
+        return None
+    if value in {"1", "true", "yes", "on"}:
+        return True
+    if value in {"0", "false", "no", "off"}:
+        return False
+    raise ValueError("supported must be true, false, or all")
+
+
+def parse_local_filter(value):
+    """Parse the model catalog local-file filter.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Parsed value."""
+    if value is None:
+        return "all"
+    value = str(value).strip().lower()
+    if value not in {"all", "complete", "missing"}:
+        raise ValueError("local must be all, complete, or missing")
+    return value
+
+
+def parse_include_files(value):
+    """Parse whether model file metadata should be included.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Parsed value."""
+    return _bool_query(value, default=False)
+
+
+def _entry_file_specs(entry, model_dir=None):
+    """Implement the entry file specs helper.
+
+    Args:
+        entry (ModelEntry): Entry value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    specs = [("model", entry.relpath, model_path_for(entry, model_dir))]
+    config_path = config_path_for(entry, model_dir)
+    if entry.config_relpath and config_path is not None:
+        specs.append(("config", entry.config_relpath, config_path))
+    specs.extend(
+        ("auxiliary", relpath, path) for relpath, path in zip(entry.auxiliary_relpaths, auxiliary_paths_for(entry, model_dir))
+    )
+    return specs
+
+
+def local_file_status(entry, model_dir=None):
+    """Report whether a catalog model exists in the local model directory.
+
+    Args:
+        entry (ModelEntry): Entry value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    specs = _entry_file_specs(entry, model_dir)
+    missing = [relpath for _role, relpath, path in specs if not path.is_file()]
+    return {
+        "complete": not missing,
+        "missing_count": len(missing),
+    }
+
+
+def catalog_model_files(entry, model_dir=None, source="modelscope", endpoint=None):
+    """Build file metadata for a catalog model.
+
+    Args:
+        entry (ModelEntry): Entry value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+        source (str, optional): Download source name. Defaults to "modelscope".
+        endpoint (str | None, optional): Optional custom download endpoint. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    files = []
+    for role, relpath, path in _entry_file_specs(entry, model_dir):
+        exists = path.is_file()
+        try:
+            size_bytes = path.stat().st_size if exists else 0
+        except OSError:
+            exists = False
+            size_bytes = 0
+        files.append(
+            {
+                "role": role,
+                "relpath": relpath,
+                "exists": exists,
+                "size_bytes": size_bytes,
+                "remote_url": remote_url(relpath, source=source, endpoint=endpoint),
+            }
+        )
+    return files
+
+
+def catalog_model_card(entry, model_dir=None, source="modelscope", endpoint=None, include_files=False):
+    """Build a summary card for a catalog model.
+
+    Args:
+        entry (ModelEntry): Entry value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+        source (str, optional): Download source name. Defaults to "modelscope".
+        endpoint (str | None, optional): Optional custom download endpoint. Defaults to None.
+        include_files (Any, optional): Include files value. Defaults to False.
+
+    Returns:
+        Any: Computed result."""
+    category = entry.category_path or entry.primary_category
+    pymss = {
+        "name": entry.name,
+        "aliases": list(entry.aliases),
+        "model_type": entry.model_type,
+        "architecture": entry.architecture,
+        "category": category,
+        "primary_category": entry.primary_category,
+        "secondary_category": entry.secondary_category,
+        "target_stem": entry.target_stem,
+        "supported": entry.supported,
+        "unsupported_reason": entry.unsupported_reason,
+        "size_bytes": entry.size_bytes,
+        "local": local_file_status(entry, model_dir),
+        "remote": {
+            "available": True,
+            "source": source,
+            "endpoint": endpoint,
+        },
+    }
+    if include_files:
+        pymss["files"] = catalog_model_files(entry, model_dir, source=source, endpoint=endpoint)
+    return {
+        "id": entry.name,
+        "object": "pymss.model_catalog_entry",
+        "owned_by": "pymss",
+        "pymss": pymss,
+    }
+
+
+def catalog_model_detail(model, model_dir=None, source="modelscope", endpoint=None):
+    """Build the detailed response for one catalog model.
+
+    Args:
+        model (str): Model value.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+        source (str, optional): Download source name. Defaults to "modelscope".
+        endpoint (str | None, optional): Optional custom download endpoint. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    entry = get_model_entry(model)
+    return catalog_model_card(entry, model_dir=model_dir, source=source, endpoint=endpoint, include_files=True)
+
+
+def filter_catalog_models(category=None, supported=True, local="all", q=None, model_dir=None):
+    """Filter catalog models by category, support, local files, and text query.
+
+    Args:
+        category (Any, optional): Category value. Defaults to None.
+        supported (bool | None, optional): Optional support-status filter. Defaults to None.
+        local (Any, optional): Local value. Defaults to 'all'.
+        q (Any, optional): Q value. Defaults to None.
+        model_dir (str | os.PathLike | None, optional): Local model cache directory. Uses the package default when None. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    rows = list_models(category=category, supported=supported)
+    query = str(q or "").strip().lower()
+    if query:
+        rows = [
+            entry
+            for entry in rows
+            if query in entry.name.lower()
+            or any(query in alias.lower() for alias in entry.aliases)
+            or query in (entry.architecture or "").lower()
+            or query in (entry.target_stem or "").lower()
+        ]
+    if local != "all":
+        want_complete = local == "complete"
+        rows = [entry for entry in rows if local_file_status(entry, model_dir)["complete"] is want_complete]
+    return rows

tools/pymss/server/state.py

@@ -0,0 +1,385 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+
+from ..config import load_config
+from ..logger import get_separation_logger
+from ..model_download import download_model
+from ..model_registry import create_separator, resolve_model
+from ..separator import INFERENCE_PARAM_TARGETS, PASSTHROUGH_INFERENCE_PARAMS
+from .config import ServerConfig
+
+
+DEFAULT_ENDPOINT = object()
+FLOAT_INFERENCE_PARAMS = frozenset({"post_process_threshold", "overlap"})
+
+
+VR_SUPPORTED_PARAMETERS = {
+    "aggression",
+    "batch_size",
+    "enable_post_process",
+    "enable_tta",
+    "fuse_conv_bn",
+    "high_end_process",
+    "mps_model_backend",
+    "mps_model_compute_dtype",
+    "normalize",
+    "post_process_threshold",
+    "use_amp",
+    "use_channels_last",
+    "window_size",
+}
+
+
+class InferenceParameterError(ValueError):
+    """Exception raised for unsupported inference parameters."""
+
+    pass
+
+
+class RequestLimiter:
+    """Async request limiter backed by a semaphore.
+
+    Args:
+        limit (int): Limit value.
+    """
+
+    def __init__(self, limit):
+        """Initialize the instance.
+
+        Args:
+            limit (int): Limit value.
+
+        Returns:
+            None: This method completes for its side effects."""
+        self.limit = max(1, int(limit))
+        self.active = 0
+        self.lock = asyncio.Lock()
+
+    async def acquire(self):
+        """Acquire value.
+
+        Args:
+            None: This callable does not accept user-provided arguments.
+
+        Returns:
+            Any: Computed result."""
+        async with self.lock:
+            if self.active >= self.limit:
+                return False
+            self.active += 1
+            return True
+
+    async def release(self):
+        """Release value.
+
+        Args:
+            None: This callable does not accept user-provided arguments.
+
+        Returns:
+            None: This callable completes for its side effects."""
+        async with self.lock:
+            self.active = max(0, self.active - 1)
+
+
+@dataclass
+class LoadedModel:
+    """Container for one loaded separator and its metadata."""
+
+    separator: object
+    entry: object
+    resolved: dict
+    requested_model: str
+    model_id: str
+    sample_rate: int
+    instruments: tuple[str, ...]
+    device: str
+    inference_params: dict
+    supported_parameters: dict[str, list[str]]
+    audio_params: dict = field(default_factory=dict)
+
+    def is_model_id(self, model):
+        """Return whether model ID.
+
+        Args:
+            model (str): Model value.
+
+        Returns:
+            bool: True when the condition is satisfied."""
+        return str(model or "") == self.model_id
+
+
+@dataclass
+class ServerState:
+    """Mutable server state for the currently loaded model."""
+
+    config: ServerConfig
+    logger: object
+    operation_lock: asyncio.Lock
+    limiter: RequestLimiter
+    model_lock: asyncio.Lock
+    inference_lock: asyncio.Lock
+    download_lock: asyncio.Lock
+    loaded: LoadedModel | None = None
+    model_loading: bool = False
+    model_loading_target: str | None = None
+    model_downloading: bool = False
+    model_downloading_target: str | None = None
+
+    def is_loaded_model(self, model):
+        """Return whether loaded model.
+
+        Args:
+            model (str): Model value.
+
+        Returns:
+            bool: True when the condition is satisfied."""
+        return self.loaded is not None and self.loaded.is_model_id(model)
+
+
+def _section(config, section):
+    """Implement the section helper.
+
+    Args:
+        config (AttrDict | dict): Loaded pymss configuration.
+        section (Mapping | None): Section value.
+
+    Returns:
+        Any: Computed result."""
+    if config is None:
+        return None
+    if isinstance(config, dict):
+        return config.get(section)
+    return getattr(config, section, None)
+
+
+def _contains(section, key):
+    """Implement the contains helper.
+
+    Args:
+        section (Mapping | None): Section value.
+        key (str): Key value.
+
+    Returns:
+        Any: Computed result."""
+    if section is None:
+        return False
+    if isinstance(section, dict):
+        return key in section
+    return hasattr(section, key)
+
+
+def _is_parameter_supported(config, model_type, key, section_name):
+    """Return whether parameter supported.
+
+    Args:
+        config (AttrDict | dict): Loaded pymss configuration.
+        model_type (Any): Model type value.
+        key (str): Key value.
+        section_name (str): Section name value.
+
+    Returns:
+        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
+    return _contains(_section(config, section_name), config_key)
+
+
+def supported_parameters(config, model_type):
+    """Return inference parameters supported by a loaded model config.
+
+    Args:
+        config (AttrDict | dict): Loaded pymss configuration.
+        model_type (Any): Model type value.
+
+    Returns:
+        Any: Computed result."""
+    grouped: dict[str, list[str]] = {}
+    for key, section_name in INFERENCE_PARAM_TARGETS.items():
+        if not _is_parameter_supported(config, model_type, key, section_name):
+            continue
+        grouped.setdefault(section_name, []).append(key)
+    return grouped
+
+
+def validate_inference_params(params, config, model_type):
+    """Validate user-provided inference parameters for a model.
+
+    Args:
+        params (dict | None): Inference parameter overrides.
+        config (AttrDict | dict): Loaded pymss configuration.
+        model_type (Any): Model type value.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    for key, value in params.items():
+        section_name = INFERENCE_PARAM_TARGETS.get(key)
+        if section_name is None:
+            raise InferenceParameterError(f"Unknown inference parameter: {key}")
+        if not _is_parameter_supported(config, model_type, key, section_name):
+            raise InferenceParameterError(f"Inference parameter {key!r} is not supported by this model")
+        if key in PASSTHROUGH_INFERENCE_PARAMS:
+            continue
+        try:
+            float(value) if key in FLOAT_INFERENCE_PARAMS else int(value)
+        except (TypeError, ValueError):
+            raise InferenceParameterError(f"Inference parameter {key!r} must be numeric")
+
+
+def _preload_config(resolved):
+    """Implement the preload config helper.
+
+    Args:
+        resolved (Any): Resolved value.
+
+    Returns:
+        Any: Computed result."""
+    model_type = resolved["model_type"]
+    if model_type == "vr":
+        return None
+    config_path = resolved.get("config_path")
+    return load_config(config_path) if config_path else None
+
+
+def _resolve_existing_or_download(model, model_dir, source, endpoint):
+    """Resolve existing or download.
+
+    Args:
+        model (str): Model value.
+        model_dir (str | os.PathLike | None): Local model cache directory. Uses the package default when None.
+        source (str): Download source name.
+        endpoint (str | None): Optional custom download endpoint.
+
+    Returns:
+        Any: Resolved value."""
+    try:
+        return resolve_model(model, model_dir=model_dir, require_supported=True, require_exists=True)
+    except FileNotFoundError:
+        download_model(model, model_dir=model_dir, source=source, endpoint=endpoint)
+        return resolve_model(model, model_dir=model_dir, require_supported=True, require_exists=True)
+
+
+def load_model(config, model, source=None, endpoint=DEFAULT_ENDPOINT, inference_params=None):
+    """Resolve and load a model into server state.
+
+    Args:
+        config (AttrDict | dict): Loaded pymss configuration.
+        model (str): Model value.
+        source (str, optional): Download source name. Defaults to "modelscope".
+        endpoint (str | None, optional): Optional custom download endpoint. Defaults to None.
+        inference_params (dict | None, optional): Inference params value. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    source = source or config.source
+    endpoint = config.endpoint if endpoint is DEFAULT_ENDPOINT else endpoint
+    params = dict(config.inference_params or {})
+    if inference_params is not None:
+        params.update(inference_params)
+
+    resolved = _resolve_existing_or_download(model, config.model_dir, source, endpoint)
+    pre_config = _preload_config(resolved)
+    validate_inference_params(params, pre_config, resolved["model_type"])
+
+    separator = create_separator(
+        model,
+        model_dir=config.model_dir,
+        device=config.device,
+        device_ids=config.device_ids or [0],
+        output_format="wav",
+        store_dirs="results",
+        logger=get_separation_logger(),
+        debug=config.debug,
+        inference_params=params,
+    )
+    instruments = tuple(str(item) for item in separator.config.training.instruments)
+    sample_rate = int(separator.config.audio.get("sample_rate", 44100))
+    entry = resolved["entry"]
+    model_type = getattr(separator, "model_type", resolved["model_type"])
+    return LoadedModel(
+        separator=separator,
+        entry=entry,
+        resolved=resolved,
+        requested_model=model,
+        model_id=entry.name,
+        sample_rate=sample_rate,
+        instruments=instruments,
+        device=separator.device,
+        inference_params=params,
+        supported_parameters=supported_parameters(separator.config, model_type),
+        audio_params=dict(getattr(separator, "audio_params", {}) or {}),
+    )
+
+
+def close_loaded_model(loaded):
+    """Close and release resources held by a loaded model.
+
+    Args:
+        loaded (LoadedModel): Loaded value.
+
+    Returns:
+        None: This callable completes for its side effects."""
+    if loaded is None:
+        return
+    separator = loaded.separator
+    close = getattr(separator, "close", None)
+    if close is not None:
+        close()
+
+
+def load_state(config):
+    """Create the initial server state.
+
+    Args:
+        config (AttrDict | dict): Loaded pymss configuration.
+
+    Returns:
+        Any: Computed result."""
+    logger = get_separation_logger()
+    state = ServerState(
+        config=config,
+        logger=logger,
+        operation_lock=asyncio.Lock(),
+        limiter=RequestLimiter(config.max_queue_size),
+        model_lock=asyncio.Lock(),
+        inference_lock=asyncio.Lock(),
+        download_lock=asyncio.Lock(),
+    )
+    if config.model:
+        state.loaded = load_model(config, config.model)
+    return state
+
+
+def model_card(loaded):
+    """Build metadata for the currently loaded model.
+
+    Args:
+        loaded (LoadedModel): Loaded value.
+
+    Returns:
+        Any: Computed result."""
+    entry = loaded.entry
+    return {
+        "id": loaded.model_id,
+        "object": "model",
+        "created": 0,
+        "owned_by": "pymss",
+        "pymss": {
+            "catalog_name": entry.name,
+            "model_type": entry.model_type,
+            "architecture": entry.architecture,
+            "category": entry.category_path or entry.primary_category,
+            "catalog_target_stem": entry.target_stem,
+            "supported": entry.supported,
+            "sample_rate": loaded.sample_rate,
+            "instruments": list(loaded.instruments),
+            "instruments_source": "separator.config.training.instruments",
+            "supported_parameters": loaded.supported_parameters,
+        },
+    }

tools/pymss/server/webui.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from fastapi import HTTPException
+from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
+
+
+WEBUI_STATIC_DIR = Path(__file__).with_name("webui_static")
+
+
+def _missing_assets_response():
+    """Implement the missing assets response helper.
+
+    Args:
+        None: This callable does not accept user-provided arguments.
+
+    Returns:
+        Any: Computed result."""
+    return JSONResponse(
+        status_code=500,
+        content={
+            "error": {
+                "message": "WebUI static assets are not built. Build the WebUI according to https://github.com/pymss-project/pymss/blob/main/README.md, then start the server again.",
+                "type": "server_error",
+                "param": None,
+                "code": "webui_assets_missing",
+            }
+        },
+    )
+
+
+def _index_path(static_dir):
+    """Implement the index path helper.
+
+    Args:
+        static_dir (str | os.PathLike): Static dir value.
+
+    Returns:
+        Any: Computed result."""
+    return Path(static_dir) / "index.html"
+
+
+def _asset_path(static_dir, asset_path):
+    """Implement the asset path helper.
+
+    Args:
+        static_dir (str | os.PathLike): Static dir value.
+        asset_path (str): Asset path value.
+
+    Returns:
+        Any: Computed result."""
+    assets_root = (Path(static_dir) / "assets").resolve()
+    candidate = (assets_root / asset_path).resolve()
+    try:
+        candidate.relative_to(assets_root)
+    except ValueError:
+        raise HTTPException(status_code=404)
+    return candidate
+
+
+def _index_response(static_dir):
+    """Implement the index response helper.
+
+    Args:
+        static_dir (str | os.PathLike): Static dir value.
+
+    Returns:
+        Any: Computed result."""
+    index = _index_path(static_dir)
+    if not index.is_file():
+        return _missing_assets_response()
+    return FileResponse(index, media_type="text/html; charset=utf-8")
+
+
+def register_webui_routes(app, static_dir=None):
+    """Register WebUI static-file and SPA fallback routes.
+
+    Args:
+        app (FastAPI): App value.
+        static_dir (str | os.PathLike, optional): Static dir value. Defaults to None.
+
+    Returns:
+        Any: Computed result."""
+    static_dir = Path(static_dir) if static_dir is not None else WEBUI_STATIC_DIR
+
+    @app.get("/ui")
+    async def redirect_webui_root():
+        """Implement the redirect webui root helper.
+
+        Args:
+            None: This callable does not accept user-provided arguments.
+
+        Returns:
+            Any: Computed result."""
+        return RedirectResponse(url="/ui/", status_code=307)
+
+    @app.get("/ui/")
+    async def webui_index():
+        """Implement the webui index helper.
+
+        Args:
+            None: This callable does not accept user-provided arguments.
+
+        Returns:
+            Any: Computed result."""
+        return _index_response(static_dir)
+
+    @app.get("/ui/assets/{asset_path:path}")
+    async def webui_asset(asset_path: str):
+        """Implement the webui asset helper.
+
+        Args:
+            asset_path (str): Asset path value.
+
+        Returns:
+            Any: Computed result."""
+        path = _asset_path(static_dir, asset_path)
+        if not path.is_file():
+            raise HTTPException(status_code=404)
+        return FileResponse(path)
+
+    @app.get("/ui/{_spa_path:path}")
+    async def webui_spa_fallback(_spa_path: str):
+        """Implement the webui spa fallback helper.
+
+        Args:
+            _spa_path (Any):  spa path value.
+
+        Returns:
+            Any: Computed result."""
+        return _index_response(static_dir)

tools/pymss/server/webui_static/index.html

@@ -0,0 +1,13 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>pymss WebUI</title>
+    <script type="module" crossorigin src="/ui/assets/index-C2f3RgYK.js"></script>
+    <link rel="stylesheet" crossorigin href="/ui/assets/index-Dd_kG96D.css">
+  </head>
+  <body>
+    <div id="app"></div>
+  </body>
+</html>

tools/pymss/workflow.py

@@ -0,0 +1,715 @@
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable
+
+import numpy as np
+import yaml
+
+
+WORKFLOW_TEMPLATE = """version: 1
+
+defaults:
+  device: auto
+  output_format: wav
+  model_dir: null
+  inference_params:
+    normalize: false
+
+steps:
+  - id: split
+    model: bs_roformer_voc_hyperacev2
+    input: input
+    stems: [vocals, other]
+    inference_params:
+      overlap_size: 48000
+    save:
+      vocals: vocal
+      other: other
+
+  - id: dereverb
+    model: UVR-DeReverb-aufr33-jarredou_4band_v4_ms_fullband
+    input: split.other
+    stems: [Dry]
+    inference_params:
+      overlap_size: 22050
+    save:
+      Dry: dry
+
+  - id: harmony
+    model: your_harmony_model
+    input: dereverb.Dry
+    stems: [other]
+    inference_params:
+      overlap_size: 22050
+    save:
+      other: harmony_other
+"""
+
+
+_STEP_ID_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_-]*$")
+_OUTPUT_FORMATS = {"wav", "flac", "mp3", "m4a"}
+_OUTPUT_LAYOUTS = {"folders", "flat"}
+_DEFAULT_AUDIO_PARAMS = {
+    "wav_bit_depth": "FLOAT",
+    "flac_bit_depth": "PCM_24",
+    "mp3_bit_rate": "320k",
+    "m4a_bit_rate": "512k",
+    "m4a_codec": "aac",
+    "m4a_aac_at_quality": 2,
+}
+
+
+class WorkflowError(ValueError):
+    """Raised when a workflow definition or run is invalid."""
+
+
+@dataclass(frozen=True)
+class WorkflowStep:
+    id: str
+    model: str | None = None
+    input: str = "input"
+    stems: list[str] | None = None
+    save: dict[str, Any] = field(default_factory=dict)
+    model_type: str | None = None
+    model_path: str | None = None
+    config_path: str | None = None
+    device: str | None = None
+    model_dir: str | None = None
+    output_format: str | None = None
+    inference_params: dict[str, Any] = field(default_factory=dict)
+    use_tta: bool | None = None
+
+
+@dataclass(frozen=True)
+class Workflow:
+    version: int
+    defaults: dict[str, Any]
+    steps: list[WorkflowStep]
+
+
+@dataclass(frozen=True)
+class AudioArtifact:
+    audio: np.ndarray
+    sample_rate: int
+
+
+@dataclass
+class WorkflowTrackState:
+    path: str
+    track_name: str
+    artifacts: dict[str, AudioArtifact] = field(default_factory=dict)
+    active: bool = True
+
+
+def load_workflow_file(path: str | os.PathLike) -> Workflow:
+    """Load a workflow YAML/JSON file."""
+    workflow_path = Path(path)
+    try:
+        data = yaml.safe_load(workflow_path.read_text(encoding="utf-8"))
+    except yaml.YAMLError as exc:
+        raise WorkflowError(f"Invalid workflow YAML: {exc}") from exc
+    except OSError as exc:
+        raise WorkflowError(f"Cannot read workflow file: {workflow_path}") from exc
+    return load_workflow_data(data)
+
+
+def load_workflow_data(data: Any) -> Workflow:
+    """Parse workflow data from a Python mapping."""
+    if not isinstance(data, dict):
+        raise WorkflowError("Workflow file must contain a mapping.")
+    version = data.get("version")
+    if version != 1:
+        raise WorkflowError("workflow version must be 1.")
+    defaults = data.get("defaults") or {}
+    if not isinstance(defaults, dict):
+        raise WorkflowError("defaults must be a mapping.")
+    raw_steps = data.get("steps")
+    if not isinstance(raw_steps, list) or not raw_steps:
+        raise WorkflowError("steps must be a non-empty list.")
+    steps = [_parse_step(index, item) for index, item in enumerate(raw_steps, start=1)]
+    workflow = Workflow(version=int(version), defaults=dict(defaults), steps=steps)
+    validate_workflow_structure(workflow)
+    return workflow
+
+
+def write_workflow_template(path: str | os.PathLike, *, overwrite: bool = False) -> Path:
+    """Write a starter workflow YAML file."""
+    output_path = Path(path)
+    if output_path.exists() and not overwrite:
+        raise WorkflowError(f"Workflow file already exists: {output_path}")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(WORKFLOW_TEMPLATE, encoding="utf-8")
+    return output_path
+
+
+def validate_workflow(
+    workflow: Workflow,
+    *,
+    model_dir: str | os.PathLike | None = None,
+    require_model_files: bool = False,
+    model_resolver: Callable[..., Any] | None = None,
+) -> Workflow:
+    """Validate workflow references and optionally model catalog entries."""
+    validate_workflow_structure(workflow)
+    _validate_step_references(workflow)
+    if model_resolver is not None or require_model_files:
+        resolver = model_resolver or _default_model_resolver
+        for step in workflow.steps:
+            if step.model_path:
+                _validate_explicit_model_files(step, require_model_files=require_model_files)
+                continue
+            step_model_dir = _step_option(workflow, step, "model_dir", model_dir)
+            resolver(
+                step.model,
+                model_dir=step_model_dir,
+                require_supported=True,
+                require_exists=require_model_files,
+            )
+    return workflow
+
+
+def validate_workflow_structure(workflow: Workflow) -> Workflow:
+    """Validate syntax that does not require cross-step analysis."""
+    seen = set()
+    for step in workflow.steps:
+        if not _STEP_ID_RE.match(step.id):
+            raise WorkflowError(f"Invalid step id {step.id!r}.")
+        if step.id in seen:
+            raise WorkflowError(f"Duplicate step id: {step.id}")
+        seen.add(step.id)
+        if bool(step.model) == bool(step.model_path):
+            raise WorkflowError(f"step {step.id!r} requires exactly one of model or model_path.")
+        if step.model_path and not step.model_type:
+            raise WorkflowError(f"step {step.id!r} requires model_type when model_path is used.")
+        if step.stems is not None and not step.stems:
+            raise WorkflowError(f"step {step.id!r} stems must not be empty.")
+        if step.output_format is not None and str(step.output_format).lower() not in _OUTPUT_FORMATS:
+            raise WorkflowError(f"step {step.id!r} has unsupported output_format {step.output_format!r}.")
+    default_format = workflow.defaults.get("output_format")
+    if default_format is not None and str(default_format).lower() not in _OUTPUT_FORMATS:
+        raise WorkflowError(f"defaults.output_format must be one of: {sorted(_OUTPUT_FORMATS)}.")
+    default_inference_params = workflow.defaults.get("inference_params")
+    if default_inference_params is not None and not isinstance(default_inference_params, dict):
+        raise WorkflowError("defaults.inference_params must be a mapping.")
+    return workflow
+
+
+class WorkflowRunner:
+    """Run a parsed pymss workflow over one file or a direct folder."""
+
+    def __init__(
+        self,
+        workflow: Workflow,
+        *,
+        model_dir: str | os.PathLike | None = None,
+        device: str | None = None,
+        output_format: str | None = None,
+        download: bool = False,
+        source: str = "modelscope",
+        endpoint: str | None = None,
+        audio_params: dict[str, Any] | None = None,
+        logger: Any = None,
+        debug: bool = False,
+        separator_factory: Callable[..., Any] | None = None,
+        audio_loader: Callable[..., Any] | None = None,
+        audio_saver: Callable[..., Any] | None = None,
+        continue_on_error: bool = False,
+        output_layout: str = "folders",
+    ):
+        self.workflow = validate_workflow(workflow)
+        self.model_dir = model_dir
+        self.device = device
+        self.output_format = output_format
+        self.download = bool(download)
+        self.source = source
+        self.endpoint = endpoint
+        self.audio_params = {**_DEFAULT_AUDIO_PARAMS, **(audio_params or {})}
+        self.logger = logger
+        self.debug = bool(debug)
+        self.separator_factory = separator_factory or _default_separator_factory
+        self.audio_loader = audio_loader or _default_audio_loader
+        self.audio_saver = audio_saver or _default_audio_saver
+        self.continue_on_error = bool(continue_on_error)
+        self.output_layout = _validate_output_layout(output_layout)
+
+    def run(self, input_path: str | os.PathLike, output_dir: str | os.PathLike) -> list[str]:
+        """Run the workflow and return successfully processed basenames."""
+        paths = _input_files(input_path)
+        output_root = Path(output_dir)
+        tracks = [
+            track
+            for path, track_name in zip(paths, _unique_track_names(paths))
+            for track in [self._load_track(path, track_name)]
+            if track is not None
+        ]
+        for step in self.workflow.steps:
+            active_tracks = [track for track in tracks if track.active]
+            if not active_tracks:
+                break
+            try:
+                with self._open_separator(step) as separator:
+                    for track in active_tracks:
+                        self._run_step_for_track(step, separator, track, output_root)
+            except Exception as exc:
+                if not self.continue_on_error:
+                    raise
+                for track in active_tracks:
+                    self._mark_track_failed(track, exc)
+        return [os.path.basename(track.path) for track in tracks if track.active]
+
+    def _load_track(self, path: str, track_name: str) -> WorkflowTrackState | None:
+        try:
+            mix, sr = self.audio_loader(path, sr=None, mono=False)
+            return WorkflowTrackState(
+                path=path,
+                track_name=track_name,
+                artifacts={"input": AudioArtifact(_to_model_audio(mix), int(sr))},
+            )
+        except Exception as exc:
+            if self.continue_on_error and self.logger is not None:
+                self.logger.warning("Cannot process workflow track %s: %s", path, exc)
+                return None
+            raise
+
+    def _run_step_for_track(
+        self,
+        step: WorkflowStep,
+        separator: Any,
+        track: WorkflowTrackState,
+        output_root: Path,
+    ) -> None:
+        try:
+            artifact = _resolve_input_artifact(track.artifacts, step.input)
+            sample_rate = int(separator.config.audio.get("sample_rate", artifact.sample_rate))
+            model_audio = _ensure_sample_rate(_to_model_audio(artifact.audio), artifact.sample_rate, sample_rate)
+            stems = _requested_stems(step)
+            if getattr(separator, "model_type", None) == "vr":
+                results = separator.separate(model_audio, pbar=False)
+            else:
+                results = separator.separate(model_audio, pbar=False, stems=stems)
+            selected = _select_results(step, results)
+            for stem, audio in selected.items():
+                track.artifacts[f"{step.id}.{stem}"] = AudioArtifact(_to_model_audio(audio), sample_rate)
+            self._save_results(step, selected, sample_rate, output_root, track.track_name)
+            del selected, results
+        except Exception as exc:
+            if not self.continue_on_error:
+                raise
+            self._mark_track_failed(track, exc)
+
+    def _mark_track_failed(self, track: WorkflowTrackState, exc: Exception) -> None:
+        track.active = False
+        if self.logger is not None:
+            self.logger.warning("Cannot process workflow track %s: %s", track.path, exc)
+
+    def _open_separator(self, step: WorkflowStep):
+        if self.download and step.model:
+            from .model_download import download_model
+
+            download_model(
+                step.model,
+                model_dir=_step_option(self.workflow, step, "model_dir", self.model_dir),
+                source=self.source,
+                endpoint=self.endpoint,
+            )
+        separator_kwargs = {
+            "model_dir": _step_option(self.workflow, step, "model_dir", self.model_dir),
+            "device": _step_option(self.workflow, step, "device", self.device),
+            "output_format": _step_option(self.workflow, step, "output_format", self.output_format) or "wav",
+            "audio_params": self.audio_params,
+            "use_tta": bool(_step_option(self.workflow, step, "use_tta", None) or False),
+            "logger": self.logger,
+            "debug": self.debug,
+            "inference_params": _merged_inference_params(self.workflow, step),
+        }
+        model_name = step.model
+        if step.model_path:
+            model_name = Path(step.model_path).stem
+            separator_kwargs.update(
+                {
+                    "model_type": step.model_type,
+                    "model_path": step.model_path,
+                    "config_path": step.config_path,
+                }
+            )
+        separator = self.separator_factory(
+            model_name,
+            **separator_kwargs,
+        )
+        return _SeparatorContext(separator)
+
+    def _save_results(
+        self,
+        step: WorkflowStep,
+        results: dict[str, np.ndarray],
+        sample_rate: int,
+        output_root: Path,
+        track_name: str,
+    ) -> None:
+        output_format = str(_step_option(self.workflow, step, "output_format", self.output_format) or "wav").lower()
+        for stem, audio in results.items():
+            save_dirs = _save_dirs(step, stem)
+            for save_dir in save_dirs:
+                target_dir = output_root / save_dir
+                if self.output_layout == "folders":
+                    target_dir = output_root / track_name / save_dir
+                target_dir.mkdir(parents=True, exist_ok=True)
+                safe_stem = _safe_filename_part(stem)
+                target = target_dir / f"{track_name}_{safe_stem}.{output_format}"
+                self.audio_saver(str(target), _to_save_audio(audio), sample_rate, output_format, self.audio_params)
+
+
+class _SeparatorContext:
+    def __init__(self, separator):
+        self.separator = separator
+
+    def __enter__(self):
+        enter = getattr(self.separator, "__enter__", None)
+        return enter() if enter is not None else self.separator
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        exit_method = getattr(self.separator, "__exit__", None)
+        if exit_method is not None:
+            return exit_method(exc_type, exc_value, traceback)
+        close = getattr(self.separator, "close", None)
+        if close is not None:
+            close()
+        return False
+
+
+def run_workflow_file(
+    config_path: str | os.PathLike,
+    input_path: str | os.PathLike,
+    output_dir: str | os.PathLike,
+    **runner_kwargs,
+) -> list[str]:
+    """Load and run a workflow file."""
+    workflow = load_workflow_file(config_path)
+    return WorkflowRunner(workflow, **runner_kwargs).run(input_path, output_dir)
+
+
+def _validate_output_layout(value: str) -> str:
+    layout = str(value).strip().lower()
+    if layout not in _OUTPUT_LAYOUTS:
+        raise WorkflowError(f"output_layout must be one of: {sorted(_OUTPUT_LAYOUTS)}.")
+    return layout
+
+
+def _parse_step(index: int, data: Any) -> WorkflowStep:
+    if not isinstance(data, dict):
+        raise WorkflowError(f"step #{index} must be a mapping.")
+    step_id = data.get("id")
+    if not isinstance(step_id, str) or not step_id.strip():
+        raise WorkflowError(f"step #{index} requires a non-empty id.")
+    model = _parse_optional_string(data.get("model"))
+    model_path = _parse_optional_string(data.get("model_path"))
+    return WorkflowStep(
+        id=step_id.strip(),
+        model=model,
+        input=_parse_input_value(data.get("input", "input"), step_id),
+        stems=_parse_stems(data.get("stems"), step_id),
+        save=_parse_save(data.get("save"), step_id),
+        model_type=_parse_optional_string(data.get("model_type")),
+        model_path=model_path,
+        config_path=_parse_optional_string(data.get("config_path")),
+        device=_parse_optional_string(data.get("device")),
+        model_dir=_parse_optional_string(data.get("model_dir")),
+        output_format=_parse_optional_string(data.get("output_format")),
+        inference_params=_parse_mapping(data.get("inference_params"), step_id, "inference_params"),
+        use_tta=_parse_optional_bool(data.get("use_tta"), step_id, "use_tta"),
+    )
+
+
+def _parse_input_value(value: Any, step_id: str) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise WorkflowError(f"step {step_id!r} input must be a non-empty string.")
+    return value.strip()
+
+
+def _parse_stems(value: Any, step_id: str) -> list[str] | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        stems = [value]
+    elif isinstance(value, list):
+        stems = value
+    else:
+        raise WorkflowError(f"step {step_id!r} stems must be a string or list.")
+    result = [str(item).strip() for item in stems if str(item).strip()]
+    if not result:
+        raise WorkflowError(f"step {step_id!r} stems must not be empty.")
+    return result
+
+
+def _parse_save(value: Any, step_id: str) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if not isinstance(value, dict):
+        raise WorkflowError(f"step {step_id!r} save must be a mapping.")
+    result = {}
+    for stem, target in value.items():
+        stem_name = str(stem).strip()
+        if not stem_name:
+            raise WorkflowError(f"step {step_id!r} save contains an empty stem name.")
+        result[stem_name] = target
+    return result
+
+
+def _parse_mapping(value: Any, step_id: str, field_name: str) -> dict[str, Any]:
+    if value is None:
+        return {}
+    if not isinstance(value, dict):
+        raise WorkflowError(f"step {step_id!r} {field_name} must be a mapping.")
+    return dict(value)
+
+
+def _parse_optional_string(value: Any) -> str | None:
+    if value is None:
+        return None
+    value = str(value).strip()
+    return value or None
+
+
+def _parse_optional_bool(value: Any, step_id: str, field_name: str) -> bool | None:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return value
+    raise WorkflowError(f"step {step_id!r} {field_name} must be a boolean.")
+
+
+def _validate_step_references(workflow: Workflow) -> None:
+    seen = {"input"}
+    required_outputs: dict[str, set[str]] = {step.id: set() for step in workflow.steps}
+    for step in workflow.steps:
+        if step.input != "input":
+            ref_step, ref_stem = _split_artifact_ref(step.input, step.id)
+            if ref_step not in seen:
+                raise WorkflowError(f"step {step.id!r} input references unknown step: {ref_step}")
+            required_outputs.setdefault(ref_step, set()).add(ref_stem)
+        for stem in step.save:
+            required_outputs[step.id].add(stem)
+        seen.add(step.id)
+
+    for step in workflow.steps:
+        if step.stems is None:
+            continue
+        requested = {stem.lower() for stem in step.stems}
+        for stem in required_outputs.get(step.id, set()):
+            if stem.lower() not in requested:
+                raise WorkflowError(
+                    f"step {step.id!r} must request {step.id}.{stem}; add {stem!r} to stems or omit stems."
+                )
+
+
+def _split_artifact_ref(value: str, current_step_id: str) -> tuple[str, str]:
+    if "." not in value:
+        raise WorkflowError(f"step {current_step_id!r} input must be 'input' or '<step>.<stem>'.")
+    step_id, stem = value.split(".", 1)
+    step_id = step_id.strip()
+    stem = stem.strip()
+    if not step_id or not stem:
+        raise WorkflowError(f"step {current_step_id!r} input must be 'input' or '<step>.<stem>'.")
+    return step_id, stem
+
+
+def _validate_explicit_model_files(step: WorkflowStep, *, require_model_files: bool) -> None:
+    if not require_model_files:
+        return
+    missing = []
+    if step.model_path and not Path(step.model_path).is_file():
+        missing.append(step.model_path)
+    if step.config_path and not Path(step.config_path).is_file():
+        missing.append(step.config_path)
+    if missing:
+        raise FileNotFoundError("Missing model file(s): " + ", ".join(missing))
+
+
+def _resolve_input_artifact(artifacts: dict[str, AudioArtifact], ref: str) -> AudioArtifact:
+    if ref == "input":
+        return artifacts["input"]
+    if ref in artifacts:
+        return artifacts[ref]
+    ref_step, ref_stem = ref.split(".", 1)
+    for key, artifact in artifacts.items():
+        if not key.startswith(f"{ref_step}."):
+            continue
+        _, stem = key.split(".", 1)
+        if stem.lower() == ref_stem.lower():
+            return artifact
+    raise WorkflowError(f"Missing workflow input artifact: {ref}")
+
+
+def _requested_stems(step: WorkflowStep) -> list[str] | None:
+    if step.stems is not None:
+        return list(step.stems)
+    if step.save:
+        return list(step.save)
+    return None
+
+
+def _select_results(step: WorkflowStep, results: dict[str, Any]) -> dict[str, np.ndarray]:
+    requested = _requested_stems(step)
+    if requested is None:
+        requested = list(results)
+    selected = {}
+    for stem in requested:
+        actual = _find_stem(results, stem)
+        selected[actual] = np.asarray(results[actual], dtype=np.float32)
+    return selected
+
+
+def _find_stem(results: dict[str, Any], stem: str) -> str:
+    if stem in results:
+        return stem
+    lower = str(stem).lower()
+    for key in results:
+        if str(key).lower() == lower:
+            return key
+    raise WorkflowError(f"Model did not return requested stem {stem!r}. Available stems: {list(results)}")
+
+
+def _save_dirs(step: WorkflowStep, stem: str) -> list[str]:
+    if not step.save:
+        return []
+    target = _case_insensitive_get(step.save, stem)
+    if target in (None, False, ""):
+        return []
+    if target is True:
+        return [step.id]
+    if isinstance(target, list):
+        return [str(item).strip() for item in target if str(item).strip()]
+    target = str(target).strip()
+    return [target] if target else []
+
+
+def _case_insensitive_get(mapping: dict[str, Any], key: str) -> Any:
+    if key in mapping:
+        return mapping[key]
+    lower = str(key).lower()
+    for item_key, value in mapping.items():
+        if str(item_key).lower() == lower:
+            return value
+    return None
+
+
+def _to_model_audio(audio: Any) -> np.ndarray:
+    array = np.asarray(audio, dtype=np.float32)
+    if array.ndim == 1:
+        return np.ascontiguousarray(array)
+    if array.ndim != 2:
+        raise WorkflowError(f"Expected mono or stereo audio, got shape {array.shape}.")
+    if array.shape[0] in (1, 2):
+        return np.ascontiguousarray(array)
+    if array.shape[1] in (1, 2):
+        return np.ascontiguousarray(array.T)
+    raise WorkflowError(f"Expected mono or stereo audio, got shape {array.shape}.")
+
+
+def _to_save_audio(audio: Any) -> np.ndarray:
+    array = np.asarray(audio, dtype=np.float32)
+    if array.ndim == 1:
+        return np.ascontiguousarray(array)
+    if array.ndim != 2:
+        raise WorkflowError(f"Expected mono or stereo audio, got shape {array.shape}.")
+    if array.shape[1] in (1, 2):
+        return np.ascontiguousarray(array)
+    if array.shape[0] in (1, 2):
+        return np.ascontiguousarray(array.T)
+    raise WorkflowError(f"Expected mono or stereo audio, got shape {array.shape}.")
+
+
+def _ensure_sample_rate(audio: np.ndarray, current_sr: int, target_sr: int) -> np.ndarray:
+    if int(current_sr) == int(target_sr):
+        return audio
+    import librosa
+
+    return np.ascontiguousarray(
+        librosa.resample(np.asarray(audio, dtype=np.float32), orig_sr=int(current_sr), target_sr=int(target_sr), axis=-1)
+    )
+
+
+def _input_files(input_path: str | os.PathLike) -> list[str]:
+    path = Path(input_path)
+    if path.is_file():
+        return [str(path)]
+    if path.is_dir():
+        return [str(item) for item in sorted(path.iterdir()) if item.is_file()]
+    raise WorkflowError(f"Input path does not exist: {path}")
+
+
+def _unique_track_names(paths: list[str]) -> list[str]:
+    original_stems = {Path(path).stem for path in paths}
+    next_suffix: dict[str, int] = {}
+    used: set[str] = set()
+    names = []
+    for path in paths:
+        stem = Path(path).stem
+        if stem not in used:
+            used.add(stem)
+            names.append(stem)
+            continue
+        suffix = next_suffix.get(stem, 2)
+        candidate = f"{stem}_{suffix}"
+        while candidate in used or candidate in original_stems:
+            suffix += 1
+            candidate = f"{stem}_{suffix}"
+        next_suffix[stem] = suffix + 1
+        used.add(candidate)
+        names.append(candidate)
+    return names
+
+
+def _step_option(workflow: Workflow, step: WorkflowStep, key: str, override: Any = None) -> Any:
+    value = getattr(step, key, None)
+    if value is not None:
+        return value
+    if override is not None:
+        return override
+    return workflow.defaults.get(key)
+
+
+def _merged_inference_params(workflow: Workflow, step: WorkflowStep) -> dict[str, Any]:
+    defaults = workflow.defaults.get("inference_params") or {}
+    return {**defaults, **(step.inference_params or {})}
+
+
+def _safe_filename_part(value: str) -> str:
+    safe = re.sub(r"[\\/:\0]+", "_", str(value)).strip()
+    return safe or "stem"
+
+
+def _default_separator_factory(model_name: str, **kwargs):
+    model_type = kwargs.pop("model_type", None)
+    model_path = kwargs.pop("model_path", None)
+    config_path = kwargs.pop("config_path", None)
+    if model_path:
+        kwargs.pop("model_dir", None)
+        from .separator import MSSeparator
+
+        return MSSeparator(model_type=model_type, model_path=model_path, config_path=config_path, **kwargs)
+
+    from .model_registry import create_separator
+    return create_separator(model_name, **kwargs)
+
+
+def _default_model_resolver(*args, **kwargs):
+    from .model_registry import resolve_model
+
+    return resolve_model(*args, **kwargs)
+
+
+def _default_audio_loader(*args, **kwargs):
+    from .audio_io import load_audio
+
+    return load_audio(*args, **kwargs)
+
+
+def _default_audio_saver(*args, **kwargs):
+    from .audio_io import save_audio
+
+    return save_audio(*args, **kwargs)

tools/pymss_core/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 KitsuneX07
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tools/pymss_core/__init__.py

@@ -0,0 +1,25 @@
+"""Core model, configuration, and checkpoint API for music source separation.
+
+`pymss_core` contains the shared pieces used by higher-level packages:
+configuration loading, model construction, model definitions, and
+checkpoint/state-dict helpers. It intentionally does not provide file audio
+I/O, inference DSP pipelines, chunked demixing, model catalog downloads, CLI,
+HTTP server, or WebUI functionality.
+"""
+
+from .checkpoint import load_checkpoint, load_model_weights, load_state_dict, unwrap_state_dict
+from .config import AttrDict, ConfigLoader, load_config, to_attrdict, to_plain
+from .utils import get_model_from_config
+
+__all__ = (
+    "AttrDict",
+    "ConfigLoader",
+    "get_model_from_config",
+    "load_checkpoint",
+    "load_config",
+    "load_model_weights",
+    "load_state_dict",
+    "to_attrdict",
+    "to_plain",
+    "unwrap_state_dict",
+)

tools/pymss_core/checkpoint.py

@@ -0,0 +1,127 @@
+"""Checkpoint helpers shared by inference and training frontends."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+import torch
+
+
+STATE_DICT_KEYS = ("state", "state_dict", "model_state_dict")
+
+
+def unwrap_state_dict(checkpoint: Any) -> Any:
+    """Return the model state dict from common MSS checkpoint containers."""
+    if isinstance(checkpoint, dict):
+        for key in STATE_DICT_KEYS:
+            if key in checkpoint:
+                return checkpoint[key]
+    return checkpoint
+
+
+def _install_demucs_pickle_stubs() -> dict[str, ModuleType | None]:
+    import sys
+    import types
+
+    module_names = ("demucs", "demucs.demucs", "demucs.hdemucs", "demucs.htdemucs")
+    previous = {name: sys.modules.get(name) for name in module_names}
+    package = sys.modules.setdefault("demucs", types.ModuleType("demucs"))
+    package.__path__ = []
+    for module_name, class_names in {
+        "demucs": ("Demucs",),
+        "hdemucs": ("HDemucs", "HTDemucs"),
+        "htdemucs": ("HTDemucs",),
+    }.items():
+        full_name = f"demucs.{module_name}"
+        module = sys.modules.setdefault(full_name, types.ModuleType(full_name))
+        setattr(package, module_name, module)
+        for class_name in class_names:
+            if not hasattr(module, class_name):
+                setattr(module, class_name, type(class_name, (), {"__module__": full_name}))
+    return previous
+
+
+def _restore_modules(previous: dict[str, ModuleType | None]) -> None:
+    import sys
+
+    for name, module in previous.items():
+        if module is None:
+            sys.modules.pop(name, None)
+        else:
+            sys.modules[name] = module
+
+
+def _torch_load(path: str | Path, *, map_location="cpu", weights_only: bool | None = None, mmap: bool = True) -> Any:
+    kwargs: dict[str, Any] = {"map_location": map_location}
+    if weights_only is not None:
+        kwargs["weights_only"] = weights_only
+    if mmap:
+        kwargs["mmap"] = True
+    try:
+        return torch.load(path, **kwargs)
+    except TypeError:
+        kwargs.pop("mmap", None)
+        try:
+            return torch.load(path, **kwargs)
+        except TypeError:
+            kwargs.pop("weights_only", None)
+            return torch.load(path, **kwargs)
+
+
+def load_checkpoint(
+    path: str | Path,
+    *,
+    model_type: str | None = None,
+    map_location: str | torch.device = "cpu",
+    weights_only: bool | None = None,
+    mmap: bool = True,
+) -> Any:
+    """Load a checkpoint package with compatibility for common MSS formats."""
+    model_type = (model_type or "").lower()
+    if model_type in {"htdemucs", "demucs", "legacy_demucs", "legacy_tasnet"}:
+        previous = _install_demucs_pickle_stubs()
+        try:
+            return _torch_load(path, map_location=map_location, weights_only=False, mmap=mmap)
+        finally:
+            _restore_modules(previous)
+    if model_type == "apollo":
+        weights_only = False if weights_only is None else weights_only
+    return _torch_load(path, map_location=map_location, weights_only=weights_only, mmap=mmap)
+
+
+def load_state_dict(
+    path: str | Path,
+    *,
+    model_type: str | None = None,
+    map_location: str | torch.device = "cpu",
+    weights_only: bool | None = None,
+    mmap: bool = True,
+) -> Any:
+    """Load and unwrap the model state dict from a checkpoint file."""
+    return unwrap_state_dict(
+        load_checkpoint(
+            path,
+            model_type=model_type,
+            map_location=map_location,
+            weights_only=weights_only,
+            mmap=mmap,
+        )
+    )
+
+
+def load_model_weights(
+    model: torch.nn.Module,
+    checkpoint_or_path: Any,
+    *,
+    model_type: str | None = None,
+    strict: bool = True,
+    map_location: str | torch.device = "cpu",
+) -> Any:
+    """Load weights from a checkpoint package or file into a model."""
+    if isinstance(checkpoint_or_path, (str, Path)):
+        state_dict = load_state_dict(checkpoint_or_path, model_type=model_type, map_location=map_location)
+    else:
+        state_dict = unwrap_state_dict(checkpoint_or_path)
+    return model.load_state_dict(state_dict, strict=strict)

tools/pymss_core/config.py

@@ -0,0 +1,135 @@
+import re
+
+import yaml
+
+
+class ConfigLoader(yaml.FullLoader):
+    """YAML loader used by pymss-core model configuration files."""
+
+    pass
+
+
+ConfigLoader.add_implicit_resolver(
+    "tag:yaml.org,2002:float",
+    re.compile(
+        r"""^[-+]?(
+            ([0-9][0-9_]*)?\.[0-9_]+([eE][-+]?[0-9]+)?
+            |[0-9][0-9_]*[eE][-+]?[0-9]+
+            |\.(inf|Inf|INF)
+            |\.(nan|NaN|NAN)
+        )$""",
+        re.X,
+    ),
+    list("-+0123456789."),
+)
+
+
+class AttrDict(dict):
+    """Dictionary that recursively exposes keys as attributes.
+
+    Args:
+        data (Mapping | None, optional): Data value. Defaults to None.
+        **kwargs: Additional keyword arguments.
+
+    Example:
+        >>> cfg = AttrDict({"audio": {"chunk_size": 485100}})
+        >>> cfg.audio.chunk_size
+        485100"""
+
+    def __init__(self, data=None, **kwargs):
+        """Initialize the instance.
+
+        Args:
+            data (Mapping | None, optional): Data value. Defaults to None.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            None: This method completes for its side effects."""
+        super().__init__()
+        for key, value in dict(data or {}, **kwargs).items():
+            self[key] = value
+
+    def __getattr__(self, key):
+        """Return a missing attribute from the underlying mapping.
+
+        Args:
+            key (str): Key value.
+
+        Returns:
+            Any: Computed result."""
+        try:
+            return self[key]
+        except KeyError as exc:
+            raise AttributeError(key) from exc
+
+    def __setattr__(self, key, value):
+        """Store an attribute assignment in the underlying mapping.
+
+        Args:
+            key (str): Key value.
+            value (Any): Value value.
+
+        Returns:
+            None: This method completes for its side effects."""
+        self[key] = value
+
+    def __setitem__(self, key, value):
+        """Store an item after recursively converting nested dictionaries.
+
+        Args:
+            key (str): Key value.
+            value (Any): Value value.
+
+        Returns:
+            None: This method completes for its side effects."""
+        super().__setitem__(key, to_attrdict(value))
+
+
+def to_attrdict(value):
+    """Recursively convert dictionaries to AttrDict objects.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Converted value with nested dictionaries wrapped as AttrDict."""
+    if isinstance(value, dict):
+        return value if isinstance(value, AttrDict) else AttrDict(value)
+    if isinstance(value, list):
+        return [to_attrdict(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(to_attrdict(item) for item in value)
+    return value
+
+
+def to_plain(value):
+    """Recursively convert AttrDict objects back to plain Python containers.
+
+    Args:
+        value (Any): Value value.
+
+    Returns:
+        Any: Converted value using plain dictionaries, lists, and tuples."""
+    if isinstance(value, AttrDict):
+        return {key: to_plain(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [to_plain(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(to_plain(item) for item in value)
+    return value
+
+
+def load_config(path):
+    """Load a YAML model configuration file.
+
+    Args:
+        path (str | os.PathLike): File system path.
+
+    Returns:
+        AttrDict: Parsed configuration with attribute access.
+
+    Example:
+        >>> config = load_config("config.yaml")
+        >>> config.inference.batch_size"""
+    with open(path, encoding="utf-8") as f:
+        return to_attrdict(yaml.load(f, Loader=ConfigLoader))

tools/pymss_core/modules/_dsp.py

@@ -0,0 +1,84 @@
+"""Small DSP helpers needed by model definitions."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def hz_to_midi(hz):
+    """Convert frequencies in Hz to MIDI note numbers."""
+    hz = np.asarray(hz)
+    return 69.0 + 12.0 * np.log2(hz / 440.0)
+
+
+def midi_to_hz(midi):
+    """Convert MIDI note numbers to frequencies in Hz."""
+    midi = np.asarray(midi)
+    return 440.0 * np.power(2.0, (midi - 69.0) / 12.0)
+
+
+def _hz_to_mel(frequencies, *, htk=False):
+    frequencies = np.asarray(frequencies, dtype=np.float64)
+    if htk:
+        return 2595.0 * np.log10(1.0 + frequencies / 700.0)
+
+    f_sp = 200.0 / 3
+    mels = frequencies / f_sp
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    logstep = np.log(6.4) / 27.0
+    log_t = frequencies >= min_log_hz
+    mels = np.array(mels, copy=True)
+    mels[log_t] = min_log_mel + np.log(frequencies[log_t] / min_log_hz) / logstep
+    return mels
+
+
+def _mel_to_hz(mels, *, htk=False):
+    mels = np.asarray(mels, dtype=np.float64)
+    if htk:
+        return 700.0 * (np.power(10.0, mels / 2595.0) - 1.0)
+
+    f_sp = 200.0 / 3
+    freqs = f_sp * mels
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    logstep = np.log(6.4) / 27.0
+    log_t = mels >= min_log_mel
+    freqs = np.array(freqs, copy=True)
+    freqs[log_t] = min_log_hz * np.exp(logstep * (mels[log_t] - min_log_mel))
+    return freqs
+
+
+def mel_frequencies(n_mels, *, fmin=0.0, fmax=11025.0, htk=False):
+    """Return center frequencies on the mel scale, including endpoints."""
+    min_mel = _hz_to_mel(fmin, htk=htk)
+    max_mel = _hz_to_mel(fmax, htk=htk)
+    return _mel_to_hz(np.linspace(min_mel, max_mel, int(n_mels)), htk=htk)
+
+
+def fft_frequencies(*, sr, n_fft):
+    """Return FFT bin center frequencies."""
+    return np.linspace(0.0, float(sr) / 2.0, int(1 + n_fft // 2), endpoint=True)
+
+
+def mel_filterbank(sr, n_fft, n_mels=128, fmin=0.0, fmax=None, htk=False, norm="slaney", dtype=np.float32):
+    """Create a triangular mel filterbank for model initialization."""
+    if fmax is None:
+        fmax = float(sr) / 2.0
+
+    mel_f = mel_frequencies(int(n_mels) + 2, fmin=fmin, fmax=fmax, htk=htk)
+    fft_f = fft_frequencies(sr=sr, n_fft=n_fft)
+
+    fdiff = np.diff(mel_f)
+    ramps = np.subtract.outer(mel_f, fft_f)
+    lower = -ramps[:-2] / fdiff[:-1, np.newaxis]
+    upper = ramps[2:] / fdiff[1:, np.newaxis]
+    weights = np.maximum(0.0, np.minimum(lower, upper))
+
+    if norm == "slaney":
+        enorm = 2.0 / (mel_f[2 : int(n_mels) + 2] - mel_f[: int(n_mels)])
+        weights *= enorm[:, np.newaxis]
+    elif norm is not None:
+        raise ValueError(f"Unsupported mel filterbank norm: {norm!r}")
+
+    return weights.astype(dtype, copy=False)

tools/pymss_core/modules/apollo_mlx.py

@@ -0,0 +1,242 @@
+import torch
+
+from .bs_roformer.mlx_attention import _mlx_dtype, _torch_to_mlx_array, mlx_to_torch_mps
+from .look2hear.apollo import BSNet, ConvActNorm1d, ICB, RMSNorm
+
+
+def torch_to_mlx_input(tensor, dtype):
+    import mlx.core as mx
+
+    return mx.array(tensor.detach().to(dtype=dtype).cpu().numpy())
+
+
+def _mlx_param(module, name, tensor, dtype):
+    cache = getattr(module, "_pymss_mlx_full_param_cache", None)
+    if cache is None:
+        cache = {}
+        module._pymss_mlx_full_param_cache = cache
+    key = (name, tensor.data_ptr(), tensor._version, tuple(tensor.shape), dtype)
+    cached = cache.get(name)
+    if cached is not None and cached[0] == key:
+        return cached[1]
+    value = _torch_to_mlx_array(tensor, dtype)
+    cache[name] = (key, value)
+    return value
+
+
+def _reflect_pad_last(x, pad):
+    import mlx.core as mx
+
+    if pad <= 0:
+        return x
+    if x.shape[-1] <= pad:
+        raise ValueError("reflect padding requires input length greater than padding")
+    return mx.concatenate((x[..., 1 : pad + 1][..., ::-1], x, x[..., -pad - 1 : -1][..., ::-1]), axis=-1)
+
+
+def _stft(module, raw_audio, dtype):
+    import mlx.core as mx
+
+    batch_channels, length = raw_audio.shape
+    n_fft = module.win
+    hop = module.stride
+    x = _reflect_pad_last(raw_audio.astype(dtype), n_fft // 2)
+    frames = 1 + (x.shape[-1] - n_fft) // hop
+    framed = mx.as_strided(x, shape=(batch_channels, frames, n_fft), strides=(x.shape[-1], hop, 1))
+    window = _torch_to_mlx_array(module.window, torch.float32).astype(dtype)
+    spec = mx.fft.rfft(framed * window, n=n_fft, axis=-1)
+    return mx.moveaxis(spec, -1, -2), {"length": length, "n_fft": n_fft, "hop": hop, "window": window, "dtype": dtype}
+
+
+def _istft(spec, context):
+    import mlx.core as mx
+
+    n_fft = context["n_fft"]
+    hop = context["hop"]
+    frames = mx.fft.irfft(mx.moveaxis(spec, -2, -1), n=n_fft, axis=-1).astype(context["dtype"]) * context["window"]
+    frame_count = frames.shape[1]
+    full_length = n_fft + hop * (frame_count - 1)
+    positions = mx.arange(n_fft)[None, :] + hop * mx.arange(frame_count)[:, None]
+    audio = mx.zeros((frames.shape[0], full_length), dtype=context["dtype"]).at[:, positions].add(frames)
+    denom_frames = mx.broadcast_to(mx.square(context["window"])[None, :], (frame_count, n_fft))
+    denom = mx.zeros((full_length,), dtype=context["dtype"]).at[positions].add(denom_frames)
+    audio = audio / mx.maximum(denom[None, :], mx.array(1e-11, dtype=context["dtype"]))
+    pad = n_fft // 2
+    return audio[..., pad : pad + context["length"]]
+
+
+def _conv1d_ncl(conv, x, dtype):
+    import mlx.core as mx
+
+    weight = _mlx_param(conv, "weight", conv.weight, dtype).transpose(0, 2, 1)
+    y = mx.conv1d(
+        x.transpose(0, 2, 1),
+        weight,
+        stride=conv.stride[0],
+        padding=conv.padding[0],
+        dilation=conv.dilation[0],
+        groups=conv.groups,
+    )
+    if conv.bias is not None:
+        y = y + _mlx_param(conv, "bias", conv.bias, dtype)
+    return y.transpose(0, 2, 1)
+
+
+def _rms_norm(module, x, dtype):
+    import mlx.core as mx
+
+    batch, channels, frames = x.shape
+    groups = int(module.groups)
+    y = x.astype(mx.float32).reshape(batch, groups, channels // groups, frames)
+    y = y * mx.rsqrt(mx.mean(mx.square(y), axis=2, keepdims=True) + module.eps)
+    y = y.reshape(batch, channels, frames).astype(x.dtype)
+    return y * _mlx_param(module, "weight", module.weight, dtype).reshape(1, -1, 1)
+
+
+def _silu(x):
+    import mlx.core as mx
+
+    return x * mx.sigmoid(x)
+
+
+def _glu_channel(x):
+    import mlx.core as mx
+
+    a, b = mx.split(x, 2, axis=1)
+    return a * mx.sigmoid(b)
+
+
+def _module_forward(module, x, dtype):
+    if isinstance(module, torch.nn.Sequential):
+        for child in module:
+            x = _module_forward(child, x, dtype)
+        return x
+    if isinstance(module, torch.nn.Conv1d):
+        return _conv1d_ncl(module, x, dtype)
+    if isinstance(module, RMSNorm):
+        return _rms_norm(module, x, dtype)
+    if isinstance(module, torch.nn.SiLU):
+        return _silu(x)
+    if isinstance(module, torch.nn.GLU):
+        return _glu_channel(x)
+    if isinstance(module, ConvActNorm1d):
+        return _conv_act_norm(module, x, dtype)
+    if isinstance(module, ICB):
+        return _module_forward(module.blocks, x, dtype)
+    if isinstance(module, BSNet):
+        return _bsnet(module, x, dtype)
+    raise TypeError(f"unsupported Apollo layer for MLX full backend: {type(module).__name__}")
+
+
+def _conv_act_norm(module, x, dtype):
+    y = _conv1d_ncl(module.conv[0], x, dtype)
+    y = _rms_norm(module.conv[1], y, dtype)
+    y = _conv1d_ncl(module.conv[2], y, dtype)
+    y = _silu(y)
+    y = _conv1d_ncl(module.conv[4], y, dtype)
+    if module.causal:
+        y = y[..., : -module.kernel + 1]
+    return x + y
+
+
+def _apply_rope(module, x, dtype):
+    import mlx.core as mx
+
+    seq_len = x.shape[-2]
+    cos = _torch_to_mlx_array(module.cos_freq[:seq_len], dtype).reshape(1, 1, seq_len, -1)
+    sin = _torch_to_mlx_array(module.sin_freq[:seq_len], dtype).reshape(1, 1, seq_len, -1)
+    even, odd = x[..., 0::2], x[..., 1::2]
+    cos_e = cos[..., 0::2]
+    sin_e = sin[..., 0::2]
+    out = mx.zeros_like(x)
+    out = out.at[..., 0::2].add(even * cos_e - odd * sin_e)
+    out = out.at[..., 1::2].add(odd * cos_e + even * sin_e)
+    return out
+
+
+def _roformer(module, x, dtype):
+    import mlx.core as mx
+
+    batch, _, frames = x.shape
+    x_norm = _rms_norm(module.input_norm, x, dtype)
+    qkv = _conv1d_ncl(module.weight, x_norm, dtype)
+    qkv = qkv.reshape(batch, module.num_head, module.hidden_size * 3, frames).transpose(0, 1, 3, 2)
+    q, k, v = mx.split(qkv, 3, axis=-1)
+    q = _apply_rope(module, q, dtype)
+    k = _apply_rope(module, k, dtype)
+    attn = mx.fast.scaled_dot_product_attention(q, k, v, scale=module.hidden_size**-0.5, mask=None)
+    out = attn.transpose(0, 1, 3, 2).reshape(batch, -1, frames)
+    out = _conv1d_ncl(module.output, out, dtype) + x
+
+    hidden = _rms_norm(module.MLP[0], out, dtype)
+    hidden = _conv1d_ncl(module.MLP[1], hidden, dtype)
+    hidden = _silu(hidden)
+    gate, z = mx.split(hidden, 2, axis=1)
+    return out + _conv1d_ncl(module.MLP_output, _silu(gate) * z, dtype)
+
+
+def _bsnet(module, x, dtype):
+    batch, bands, channels, frames = x.shape
+    band = x.transpose(0, 3, 2, 1).reshape(batch * frames, channels, bands)
+    band = _roformer(module.band_net, band, dtype)
+    band = band.reshape(batch, frames, channels, bands).transpose(0, 3, 2, 1)
+    seq = _module_forward(module.seq_net, band.reshape(batch * bands, channels, frames), dtype)
+    return seq.reshape(batch, bands, channels, frames)
+
+
+def _feature_extractor(module, raw_audio, dtype):
+    import mlx.core as mx
+
+    mx_dtype = _mlx_dtype(dtype)
+    batch, channels, samples = raw_audio.shape
+    spec, _ = _stft(module, raw_audio.reshape(batch * channels, samples), mx_dtype)
+    features = []
+    powers = []
+    band_index = 0
+    for width, bn in zip(module.band_width, module.BN):
+        sub = spec[:, band_index : band_index + width]
+        power = mx.sqrt(mx.sum(mx.square(sub.real) + mx.square(sub.imag), axis=1, keepdims=True) + module.eps)
+        norm = sub / power
+        inp = mx.concatenate((norm.real, norm.imag, mx.log(power)), axis=1)
+        features.append(_module_forward(bn, inp.astype(mx_dtype), dtype))
+        powers.append(power)
+        band_index += width
+    return mx.stack(features, axis=1), spec
+
+
+def _estimate_spec(module, feature, batch_channels, dtype):
+    import mlx.core as mx
+
+    specs = []
+    for band_feature, output, width in zip(mx.split(feature, feature.shape[1], axis=1), module.output, module.band_width):
+        band_feature = band_feature[:, 0]
+        ri = _module_forward(output, band_feature, dtype).reshape(batch_channels, 2, width, -1)
+        specs.append(ri[:, 0] + (1j * ri[:, 1]))
+    return mx.concatenate(specs, axis=1)
+
+
+def mlx_forward_apollo_mx(module, raw_audio, dtype=torch.float16):
+    if dtype not in (torch.float16, torch.float32):
+        raise TypeError("MLX full Apollo supports torch.float16 or torch.float32 compute dtype")
+    mx_dtype = _mlx_dtype(dtype)
+    raw_audio = raw_audio.astype(mx_dtype)
+    batch, channels, samples = raw_audio.shape
+    feature, _ = _feature_extractor(module, raw_audio, dtype)
+    for block in module.net:
+        feature = _bsnet(block, feature, dtype)
+    est_spec = _estimate_spec(module, feature, batch * channels, dtype)
+    return _istft(
+        est_spec,
+        {
+            "length": samples,
+            "n_fft": module.win,
+            "hop": module.stride,
+            "window": _torch_to_mlx_array(module.window, torch.float32).astype(mx_dtype),
+            "dtype": mx_dtype,
+        },
+    ).reshape(batch, channels, -1)
+
+
+def mlx_forward_apollo(module, raw_audio, dtype=torch.float16):
+    x_mx = torch_to_mlx_input(raw_audio, dtype=dtype)
+    return mlx_to_torch_mps(mlx_forward_apollo_mx(module, x_mx, dtype), raw_audio)

tools/pymss_core/modules/bandit/bandsplit.py

@@ -0,0 +1,166 @@
+from typing import List, Tuple
+
+import torch
+from torch import nn
+from torch.utils.checkpoint import checkpoint_sequential
+
+from .core.model.bsrnn.utils import (
+    band_widths_from_specs,
+    check_no_gap,
+    check_no_overlap,
+    check_nonzero_bandwidth,
+)
+
+
+class NormFC(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        bandwidth: int,
+        in_channels: int,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__()
+        self.treat_channel_as_feature = treat_channel_as_feature
+        if normalize_channel_independently:
+            raise NotImplementedError
+
+        reim = 2
+        self.norm = nn.LayerNorm(in_channels * bandwidth * reim)
+
+        fc_in = bandwidth * reim
+        if treat_channel_as_feature:
+            fc_in *= in_channels
+        else:
+            assert emb_dim % in_channels == 0
+            emb_dim = emb_dim // in_channels
+
+        self.fc = nn.Linear(fc_in, emb_dim)
+
+    def forward(self, xb):
+        batch, n_time, in_channels, ribw = xb.shape
+        xb = self.norm(xb.reshape(batch, n_time, in_channels * ribw))
+        if self.treat_channel_as_feature:
+            return self.fc(xb)
+        return self.fc(xb.reshape(batch, n_time, in_channels, ribw)).reshape(batch, n_time, -1)
+
+
+class SequentialNormFC(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        bandwidth: int,
+        in_channels: int,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__()
+        if not treat_channel_as_feature:
+            raise NotImplementedError
+        if normalize_channel_independently:
+            raise NotImplementedError
+
+        self.combined = nn.Sequential(
+            nn.LayerNorm(in_channels * bandwidth * 2),
+            nn.Linear(in_channels * bandwidth * 2, emb_dim),
+        )
+
+    def forward(self, xb):
+        return checkpoint_sequential(self.combined, 1, xb, use_reentrant=False)
+
+
+class BandSplitModuleBase(nn.Module):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        in_channels: int,
+        norm_fc_cls: type[nn.Module],
+        complex_order: str,
+        flatten_input: bool,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__()
+        check_nonzero_bandwidth(band_specs)
+        if require_no_gap:
+            check_no_gap(band_specs)
+        if require_no_overlap:
+            check_no_overlap(band_specs)
+
+        self.band_specs = band_specs
+        self.band_widths = band_widths_from_specs(band_specs)
+        self.n_bands = len(band_specs)
+        self.emb_dim = emb_dim
+        self.complex_order = complex_order
+        self.flatten_input = flatten_input
+        self.norm_fc_modules = nn.ModuleList(
+            [
+                norm_fc_cls(
+                    emb_dim=emb_dim,
+                    bandwidth=bw,
+                    in_channels=in_channels,
+                    normalize_channel_independently=normalize_channel_independently,
+                    treat_channel_as_feature=treat_channel_as_feature,
+                )
+                for bw in self.band_widths
+            ]
+        )
+
+    def _band_view(self, x):
+        xr = torch.view_as_real(x)
+        if self.complex_order == "reim_freq":
+            return xr.permute(0, 3, 1, 4, 2)
+        if self.complex_order == "freq_reim":
+            return xr.permute(0, 3, 1, 2, 4).contiguous()
+        raise ValueError(f"unsupported complex_order: {self.complex_order}")
+
+    def forward(self, x: torch.Tensor):
+        batch, in_channels, _, n_time = x.shape
+        z = torch.zeros(
+            size=(batch, self.n_bands, n_time, self.emb_dim),
+            device=x.device,
+        )
+        xr = self._band_view(x)
+
+        for i, nfm in enumerate(self.norm_fc_modules):
+            fstart, fend = self.band_specs[i]
+            if self.complex_order == "reim_freq":
+                xb = xr[..., fstart:fend].reshape(batch, n_time, in_channels, -1)
+            else:
+                xb = xr[:, :, :, fstart:fend].reshape(batch, n_time, -1)
+            z[:, i, :, :] = nfm((xb.reshape(batch, n_time, -1) if self.flatten_input else xb).contiguous())
+
+        return z
+
+
+class _ConfiguredBandSplitModule(BandSplitModuleBase):
+    norm_fc_cls: type[nn.Module]
+    complex_order: str
+    flatten_input: bool
+
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        in_channels: int,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__(
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            in_channels=in_channels,
+            norm_fc_cls=self.norm_fc_cls,
+            complex_order=self.complex_order,
+            flatten_input=self.flatten_input,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+        )

tools/pymss_core/modules/bandit/core/__init__.py

@@ -0,0 +1,9 @@
+__all__ = ("MultiMaskMultiSourceBandSplitRNNSimple",)
+
+
+def __getattr__(name):
+    if name == "MultiMaskMultiSourceBandSplitRNNSimple":
+        from .model import MultiMaskMultiSourceBandSplitRNNSimple
+
+        return MultiMaskMultiSourceBandSplitRNNSimple
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

tools/pymss_core/modules/bandit/core/model/__init__.py

@@ -0,0 +1,9 @@
+__all__ = ("MultiMaskMultiSourceBandSplitRNNSimple",)
+
+
+def __getattr__(name):
+    if name == "MultiMaskMultiSourceBandSplitRNNSimple":
+        from .bsrnn.wrapper import MultiMaskMultiSourceBandSplitRNNSimple
+
+        return MultiMaskMultiSourceBandSplitRNNSimple
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

tools/pymss_core/modules/bandit/core/model/_spectral.py

@@ -0,0 +1,95 @@
+from typing import Dict, Optional
+
+import torch
+from torch import nn
+
+
+class _TorchSpectrogram(nn.Module):
+    def __init__(
+        self,
+        n_fft,
+        win_length,
+        hop_length,
+        window_fn,
+        wkwargs,
+        normalized,
+        center,
+        pad_mode,
+        onesided,
+    ):
+        super().__init__()
+        self.n_fft = n_fft
+        self.win_length = win_length or n_fft
+        self.hop_length = hop_length
+        self.normalized = normalized
+        self.center = center
+        self.pad_mode = pad_mode
+        self.onesided = onesided
+        self.register_buffer("window", window_fn(self.win_length, **(wkwargs or {})))
+
+    def forward(self, x):
+        leading_shape = x.shape[:-1]
+        spec = torch.stft(
+            x.reshape(-1, x.shape[-1]),
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.window,
+            center=self.center,
+            pad_mode=self.pad_mode,
+            normalized=self.normalized,
+            onesided=self.onesided,
+            return_complex=True,
+        )
+        return spec.reshape(*leading_shape, *spec.shape[-2:])
+
+
+class _TorchInverseSpectrogram(_TorchSpectrogram):
+    def forward(self, x, length=None):
+        leading_shape = x.shape[:-2]
+        audio = torch.istft(
+            x.reshape(-1, *x.shape[-2:]),
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.window,
+            center=self.center,
+            normalized=self.normalized,
+            onesided=self.onesided,
+            length=length,
+            return_complex=False,
+        )
+        return audio.reshape(*leading_shape, audio.shape[-1])
+
+
+class _SpectralComponent(nn.Module):
+    def __init__(
+        self,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        assert power is None
+        window_fn = torch.__dict__[window_fn]
+        kwargs = dict(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            normalized=normalized,
+            center=center,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+        self.stft = _TorchSpectrogram(**kwargs)
+        self.istft = _TorchInverseSpectrogram(**kwargs)

tools/pymss_core/modules/bandit/core/model/bsrnn/__init__.py

@@ -0,0 +1,17 @@
+from abc import ABC
+from typing import Iterable, Mapping, Union
+
+from torch import nn
+
+
+class BandsplitCoreBase(nn.Module, ABC):
+    band_split: nn.Module
+    tf_model: nn.Module
+    mask_estim: Union[nn.Module, Mapping[str, nn.Module], Iterable[nn.Module]]
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @staticmethod
+    def mask(x, m):
+        return x * m

tools/pymss_core/modules/bandit/core/model/bsrnn/bandsplit.py

@@ -0,0 +1,32 @@
+from typing import List, Tuple
+
+from ....bandsplit import NormFC, _ConfiguredBandSplitModule
+
+
+class BandSplitModule(_ConfiguredBandSplitModule):
+    norm_fc_cls = NormFC
+    complex_order = "reim_freq"
+    flatten_input = False
+
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        in_channel: int,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+    ) -> None:
+        super().__init__(
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            in_channels=in_channel,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+        )
+
+
+__all__ = ("BandSplitModule", "NormFC")

tools/pymss_core/modules/bandit/core/model/bsrnn/core.py

@@ -0,0 +1,171 @@
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch import nn
+
+from . import BandsplitCoreBase
+from .bandsplit import BandSplitModule
+from .maskestim import MaskEstimationModule, OverlappingMaskEstimationModule
+from .tfmodel import SeqBandModellingModule
+
+__all__ = ("MultiSourceMultiMaskBandSplitCoreRNN",)
+
+
+class MultiMaskBandSplitCoreBase(BandsplitCoreBase):
+    def forward(self, x, cond=None, compute_residual: bool = True):
+        batch, in_chan, n_freq, n_time = x.shape
+        x = x.reshape(-1, 1, n_freq, n_time)
+
+        z = self.band_split(x)
+        q = self.tf_model(z)
+
+        out = {}
+        for stem, mask_estimator in self.mask_estim.items():
+            mask = mask_estimator(q, cond=cond)
+            separated = self.mask(x, mask)
+            out[stem] = separated.reshape(batch, in_chan, n_freq, n_time)
+
+        return {"spectrogram": out}
+
+    def instantiate_mask_estim(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        cond_dim: int,
+        hidden_activation: str,
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        overlapping_band: bool = False,
+        freq_weights: Optional[List[torch.Tensor]] = None,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = True,
+        mult_add_mask: bool = False,
+    ):
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+        if mult_add_mask:
+            raise NotImplementedError("Bandit mult_add_mask is not supported by the inference-only wrapper")
+
+        stems = [stem for stem in stems if stem != "mne:+"]
+        if overlapping_band:
+            assert freq_weights is not None
+            assert n_freq is not None
+            self.mask_estim = nn.ModuleDict(
+                {
+                    stem: OverlappingMaskEstimationModule(
+                        band_specs=band_specs,
+                        freq_weights=freq_weights,
+                        n_freq=n_freq,
+                        emb_dim=emb_dim,
+                        mlp_dim=mlp_dim,
+                        in_channel=in_channel,
+                        hidden_activation=hidden_activation,
+                        hidden_activation_kwargs=hidden_activation_kwargs,
+                        complex_mask=complex_mask,
+                        use_freq_weights=use_freq_weights,
+                    )
+                    for stem in stems
+                }
+            )
+        else:
+            self.mask_estim = nn.ModuleDict(
+                {
+                    stem: MaskEstimationModule(
+                        band_specs=band_specs,
+                        emb_dim=emb_dim,
+                        mlp_dim=mlp_dim,
+                        cond_dim=cond_dim,
+                        in_channel=in_channel,
+                        hidden_activation=hidden_activation,
+                        hidden_activation_kwargs=hidden_activation_kwargs,
+                        complex_mask=complex_mask,
+                    )
+                    for stem in stems
+                }
+            )
+
+    def instantiate_bandsplit(
+        self,
+        in_channel: int,
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        emb_dim: int = 128,
+    ):
+        self.band_split = BandSplitModule(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+
+
+class MultiSourceMultiMaskBandSplitCoreRNN(MultiMaskBandSplitCoreBase):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: List[Tuple[float, float]],
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        cond_dim: int = 0,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        overlapping_band: bool = False,
+        freq_weights: Optional[List[torch.Tensor]] = None,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = True,
+        mult_add_mask: bool = False,
+    ) -> None:
+        super().__init__()
+        self.instantiate_bandsplit(
+            in_channel=in_channel,
+            band_specs=band_specs,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+
+        self.tf_model = SeqBandModellingModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+        )
+
+        self.instantiate_mask_estim(
+            in_channel=in_channel,
+            stems=stems,
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=overlapping_band,
+            freq_weights=freq_weights,
+            n_freq=n_freq,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )

tools/pymss_core/modules/bandit/core/model/bsrnn/maskestim.py

@@ -0,0 +1,20 @@
+from ....maskestim import (
+    BaseNormMLP,
+    MaskEstimationModule,
+    MaskEstimationModuleBase,
+    MaskEstimationModuleSuperBase,
+    MultAddNormMLP,
+    NormMLP,
+    OverlappingMaskEstimationModule,
+)
+
+
+__all__ = (
+    "BaseNormMLP",
+    "MaskEstimationModule",
+    "MaskEstimationModuleBase",
+    "MaskEstimationModuleSuperBase",
+    "MultAddNormMLP",
+    "NormMLP",
+    "OverlappingMaskEstimationModule",
+)

tools/pymss_core/modules/bandit/core/model/bsrnn/tfmodel.py

@@ -0,0 +1,12 @@
+from ....tfmodel import (
+    ResidualRNN,
+    TimeFrequencyModellingModule,
+    _SeqBandModellingPreset,
+)
+
+
+class SeqBandModellingModule(_SeqBandModellingPreset):
+    pass
+
+
+__all__ = ("ResidualRNN", "SeqBandModellingModule", "TimeFrequencyModellingModule")

tools/pymss_core/modules/bandit/core/model/bsrnn/utils.py

@@ -0,0 +1,355 @@
+import os
+from abc import abstractmethod
+from typing import Callable
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from ....._dsp import hz_to_midi, mel_filterbank as _mel_filterbank, midi_to_hz
+
+
+def band_widths_from_specs(band_specs):
+    return [e - i for i, e in band_specs]
+
+
+def check_nonzero_bandwidth(band_specs):
+    for fstart, fend in band_specs:
+        if fend - fstart <= 0:
+            raise ValueError("Bands cannot be zero-width")
+
+
+def check_no_overlap(band_specs):
+    fend_prev = -1
+    for fstart_curr, fend_curr in band_specs:
+        if fstart_curr <= fend_prev:
+            raise ValueError("Bands cannot overlap")
+
+
+def check_no_gap(band_specs):
+    fstart, _ = band_specs[0]
+    assert fstart == 0
+
+    fend_prev = -1
+    for fstart_curr, fend_curr in band_specs:
+        if fstart_curr - fend_prev > 1:
+            raise ValueError("Bands cannot leave gap")
+        fend_prev = fend_curr
+
+
+def create_triangular_filterbank(all_freqs, f_pts):
+    f_diff = f_pts[1:] - f_pts[:-1]
+    slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)
+    down_slopes = -slopes[:, :-2] / f_diff[:-1]
+    up_slopes = slopes[:, 2:] / f_diff[1:]
+    return torch.clamp(torch.minimum(down_slopes, up_slopes), min=0.0)
+
+
+def triangular_filterbank_from_points(all_freqs, f_pts):
+    fb = create_triangular_filterbank(all_freqs, f_pts).T
+    first_active_band = torch.nonzero(torch.sum(fb, dim=-1))[0, 0]
+    fb[first_active_band, : torch.nonzero(fb[first_active_band, :])[0, 0]] = 1.0
+    return fb
+
+
+def hz_to_bark(hz):
+    return 6 * np.arcsinh(np.asarray(hz) / 600)
+
+
+def hz_to_erb(hz):
+    a = (1000 * np.log(10)) / (24.7 * 4.37)
+    return a * np.log10(1 + 0.00437 * np.asarray(hz))
+
+
+class BandsplitSpecification:
+    def __init__(self, nfft: int, fs: int) -> None:
+        self.fs = fs
+        self.nfft = nfft
+        self.nyquist = fs / 2
+        self.max_index = nfft // 2 + 1
+
+        self.split500 = self.hertz_to_index(500)
+        self.split1k = self.hertz_to_index(1000)
+        self.split2k = self.hertz_to_index(2000)
+        self.split4k = self.hertz_to_index(4000)
+        self.split8k = self.hertz_to_index(8000)
+        self.split16k = self.hertz_to_index(16000)
+        self.split20k = self.hertz_to_index(20000)
+
+        self.above20k = [(self.split20k, self.max_index)]
+        self.above16k = [(self.split16k, self.split20k)] + self.above20k
+
+    def index_to_hertz(self, index: int):
+        return index * self.fs / self.nfft
+
+    def hertz_to_index(self, hz: float, round: bool = True):
+        index = hz * self.nfft / self.fs
+        if round:
+            index = int(np.round(index))
+        return index
+
+    def get_band_specs_with_bandwidth(self, start_index, end_index, bandwidth_hz):
+        band_specs = []
+        lower = start_index
+        while lower < end_index:
+            upper = min(int(np.floor(lower + self.hertz_to_index(bandwidth_hz))), end_index)
+            band_specs.append((lower, upper))
+            lower = upper
+        return band_specs
+
+    def bands(self, *segments):
+        return sum((self.get_band_specs_with_bandwidth(start, end, bandwidth) for start, end, bandwidth in segments), [])
+
+    @abstractmethod
+    def get_band_specs(self):
+        raise NotImplementedError
+
+
+class VocalBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+        self.version = version
+
+    def get_band_specs(self):
+        return getattr(self, f"version{self.version}")()
+
+    def version1(self):
+        return self.bands((0, self.max_index, 1000))
+
+    def version2(self):
+        return self.bands((0, self.split16k, 1000), (self.split16k, self.split20k, 2000)) + self.above20k
+
+    def version3(self):
+        return self.bands((0, self.split8k, 1000), (self.split8k, self.split16k, 2000)) + self.above16k
+
+    def version4(self):
+        return (
+            self.bands((0, self.split1k, 100), (self.split1k, self.split8k, 1000), (self.split8k, self.split16k, 2000))
+            + self.above16k
+        )
+
+    def version5(self):
+        return (
+            self.bands((0, self.split1k, 100), (self.split1k, self.split16k, 1000), (self.split16k, self.split20k, 2000))
+            + self.above20k
+        )
+
+    def version6(self):
+        return (
+            self.bands(
+                (0, self.split1k, 100),
+                (self.split1k, self.split4k, 500),
+                (self.split4k, self.split8k, 1000),
+                (self.split8k, self.split16k, 2000),
+            )
+            + self.above16k
+        )
+
+    def version7(self):
+        return (
+            self.bands(
+                (0, self.split1k, 100),
+                (self.split1k, self.split4k, 250),
+                (self.split4k, self.split8k, 500),
+                (self.split8k, self.split16k, 1000),
+                (self.split16k, self.split20k, 2000),
+            )
+            + self.above20k
+        )
+
+
+class OtherBandsplitSpecification(VocalBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int) -> None:
+        super().__init__(nfft=nfft, fs=fs, version="7")
+
+
+class BassBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, version: str = "7") -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+    def get_band_specs(self):
+        return self.bands(
+            (0, self.split500, 50),
+            (self.split500, self.split1k, 100),
+            (self.split1k, self.split4k, 500),
+            (self.split4k, self.split8k, 1000),
+            (self.split8k, self.split16k, 2000),
+        ) + [(self.split16k, self.max_index)]
+
+
+class DrumBandsplitSpecification(BandsplitSpecification):
+    def __init__(self, nfft: int, fs: int) -> None:
+        super().__init__(nfft=nfft, fs=fs)
+
+    def get_band_specs(self):
+        return self.bands(
+            (0, self.split1k, 50),
+            (self.split1k, self.split2k, 100),
+            (self.split2k, self.split4k, 250),
+            (self.split4k, self.split8k, 500),
+            (self.split8k, self.split16k, 1000),
+        ) + [(self.split16k, self.max_index)]
+
+
+class PerceptualBandsplitSpecification(BandsplitSpecification):
+    def __init__(
+        self,
+        nfft: int,
+        fs: int,
+        fbank_fn: Callable[[int, int, float, float, int], torch.Tensor],
+        n_bands: int,
+        f_min: float = 0.0,
+        f_max: float = None,
+    ) -> None:
+        super().__init__(nfft=nfft, fs=fs)
+        self.n_bands = n_bands
+        if f_max is None:
+            f_max = fs / 2
+
+        self.filterbank = fbank_fn(n_bands, fs, f_min, f_max, self.max_index)
+
+        weight_per_bin = torch.sum(self.filterbank, dim=0, keepdim=True)
+        normalized_mel_fb = self.filterbank / weight_per_bin  # (n_mels, n_freqs)
+
+        freq_weights = []
+        band_specs = []
+        for i in range(self.n_bands):
+            active_bins = torch.nonzero(self.filterbank[i, :]).squeeze().tolist()
+            if isinstance(active_bins, int):
+                active_bins = (active_bins, active_bins)
+            if len(active_bins) == 0:
+                continue
+            band_specs.append((start_index := active_bins[0], end_index := active_bins[-1] + 1))
+            freq_weights.append(normalized_mel_fb[i, start_index:end_index])
+
+        self.freq_weights = freq_weights
+        self.band_specs = band_specs
+
+    def get_band_specs(self):
+        return self.band_specs
+
+    def get_freq_weights(self):
+        return self.freq_weights
+
+    def save_to_file(self, dir_path: str) -> None:
+        os.makedirs(dir_path, exist_ok=True)
+
+        import pickle
+
+        with open(os.path.join(dir_path, "mel_bandsplit_spec.pkl"), "wb") as f:
+            pickle.dump({"band_specs": self.band_specs, "freq_weights": self.freq_weights, "filterbank": self.filterbank}, f)
+
+
+def mel_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    nfft = 2 * (n_freqs - 1)
+    fb = torch.as_tensor(
+        _mel_filterbank(
+            sr=fs,
+            n_fft=nfft,
+            n_mels=n_bands,
+            fmin=f_min,
+            fmax=f_max,
+            htk=True,
+            norm=None,
+        )
+    )
+    fb[0, 0] = 1.0
+    return fb
+
+
+class MelBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None) -> None:
+        super().__init__(fbank_fn=mel_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max)
+
+
+def musical_filterbank(n_bands, fs, f_min, f_max, n_freqs, scale="constant"):
+    nfft, f_max, f_min = 2 * (n_freqs - 1), f_max or fs / 2, fs / (2 * (n_freqs - 1))
+    df, bandwidth_mult = fs / nfft, np.power(2.0, np.log2(f_max / f_min) / n_bands)
+    hz_pts = midi_to_hz(np.linspace(max(0, hz_to_midi(f_min)), hz_to_midi(f_max), n_bands))
+    low_bins, high_bins = np.floor(hz_pts / bandwidth_mult / df).astype(int), np.ceil(hz_pts * bandwidth_mult / df).astype(int)
+
+    fb = np.zeros((n_bands, n_freqs))
+
+    for i in range(n_bands):
+        fb[i, low_bins[i] : high_bins[i] + 1] = 1.0
+
+    fb[0, : low_bins[0]] = 1.0
+    fb[-1, high_bins[-1] + 1 :] = 1.0
+
+    return torch.as_tensor(fb)
+
+
+class MusicalBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None) -> None:
+        super().__init__(fbank_fn=musical_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max)
+
+
+def bark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    nfft = 2 * (n_freqs - 1)
+    f_max = f_max or fs / 2
+    centers = np.linspace(hz_to_bark(f_min), hz_to_bark(f_max), n_bands)
+    bins = np.floor((nfft + 1) * (600 * np.sinh(centers / 6) / fs)).astype(int)
+    start, end = int(bins[0]), int(bins[-1])
+    bark_bins = hz_to_bark(np.arange(start, end) * fs / (nfft + 1))
+    fb = np.zeros((n_bands, n_freqs))
+
+    for band, center in enumerate(centers):
+        diff = bark_bins - center
+        values = np.zeros_like(diff)
+        lower = (-1.3 <= diff) & (diff <= -0.5)
+        center_mask = (-0.5 < diff) & (diff < 0.5)
+        upper = (0.5 <= diff) & (diff <= 2.5)
+        values[lower] = 10 ** (2.5 * (diff[lower] + 0.5))
+        values[center_mask] = 1
+        values[upper] = 10 ** (-(diff[upper] - 0.5))
+        fb[band, start:end] = values
+
+    return torch.as_tensor(fb)
+
+
+class BarkBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None) -> None:
+        super().__init__(fbank_fn=bark_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max)
+
+
+def triangular_bark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    return triangular_filterbank_from_points(
+        torch.linspace(0, fs // 2, n_freqs),
+        600 * torch.sinh(torch.linspace(hz_to_bark(f_min), hz_to_bark(f_max), n_bands + 2) / 6),
+    )
+
+
+class TriangularBarkBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None) -> None:
+        super().__init__(fbank_fn=triangular_bark_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max)
+
+
+def minibark_filterbank(n_bands, fs, f_min, f_max, n_freqs):
+    fb = bark_filterbank(n_bands, fs, f_min, f_max, n_freqs)
+    fb[fb < np.sqrt(0.5)] = 0.0
+    return fb
+
+
+class MiniBarkBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None) -> None:
+        super().__init__(fbank_fn=minibark_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max)
+
+
+def erb_filterbank(
+    n_bands: int,
+    fs: int,
+    f_min: float,
+    f_max: float,
+    n_freqs: int,
+) -> Tensor:
+    A = (1000 * np.log(10)) / (24.7 * 4.37)
+    return triangular_filterbank_from_points(
+        torch.linspace(0, fs // 2, n_freqs),
+        (torch.pow(10, torch.linspace(hz_to_erb(f_min), hz_to_erb(f_max), n_bands + 2) / A) - 1) / 0.00437,
+    )
+
+
+class EquivalentRectangularBandsplitSpecification(PerceptualBandsplitSpecification):
+    def __init__(self, nfft: int, fs: int, n_bands: int, f_min: float = 0.0, f_max: float = None) -> None:
+        super().__init__(fbank_fn=erb_filterbank, nfft=nfft, fs=fs, n_bands=n_bands, f_min=f_min, f_max=f_max)

tools/pymss_core/modules/bandit/core/model/bsrnn/wrapper.py

@@ -0,0 +1,235 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+
+from .._spectral import _SpectralComponent
+from .core import MultiSourceMultiMaskBandSplitCoreRNN
+from .utils import (
+    BarkBandsplitSpecification,
+    EquivalentRectangularBandsplitSpecification,
+    MelBandsplitSpecification,
+    MusicalBandsplitSpecification,
+    TriangularBarkBandsplitSpecification,
+    VocalBandsplitSpecification,
+)
+
+__all__ = ("MultiMaskMultiSourceBandSplitRNNSimple",)
+
+
+def get_band_specs(band_specs, n_fft, fs, n_bands=None):
+    if not isinstance(band_specs, str):
+        return band_specs, None, False
+
+    if band_specs in ["dnr:speech", "dnr:vox7", "musdb:vocals", "musdb:vox7"]:
+        bsm = VocalBandsplitSpecification(nfft=n_fft, fs=fs).get_band_specs()
+        freq_weights = None
+        overlapping_band = False
+    elif "tribark" in band_specs:
+        assert n_bands is not None
+        specs = TriangularBarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif "bark" in band_specs:
+        assert n_bands is not None
+        specs = BarkBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif "erb" in band_specs:
+        assert n_bands is not None
+        specs = EquivalentRectangularBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif "musical" in band_specs:
+        assert n_bands is not None
+        specs = MusicalBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    elif band_specs == "dnr:mel" or "mel" in band_specs:
+        assert n_bands is not None
+        specs = MelBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+        bsm = specs.get_band_specs()
+        freq_weights = specs.get_freq_weights()
+        overlapping_band = True
+    else:
+        raise ValueError(f"Unsupported band_specs: {band_specs}")
+
+    return bsm, freq_weights, overlapping_band
+
+
+class MultiMaskMultiSourceBandSplitBaseSimple(_SpectralComponent):
+    mps_model_backend = "torch"
+    mps_model_compute_dtype = torch.float16
+
+    def __init__(
+        self,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+    ) -> None:
+        super().__init__(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+        self.band_specs, self.freq_weights, self.overlapping_band = get_band_specs(
+            band_specs,
+            n_fft,
+            fs,
+            n_bands,
+        )
+        self.stems = stems
+
+    def set_mps_model_backend(self, backend=None, compute_dtype=None):
+        backend = (backend or "torch").lower()
+        if backend not in ("torch", "mlx_full"):
+            raise ValueError("mps_model_backend must be 'torch' or 'mlx_full'")
+        self.mps_model_backend = backend
+        if compute_dtype is None:
+            return
+        if isinstance(compute_dtype, str):
+            compute_dtype = {
+                "float16": torch.float16,
+                "fp16": torch.float16,
+                "float32": torch.float32,
+                "fp32": torch.float32,
+            }.get(compute_dtype.lower(), compute_dtype)
+        if compute_dtype not in (torch.float16, torch.float32):
+            raise ValueError("mps_model_compute_dtype must be 'float16' or 'float32'")
+        self.mps_model_compute_dtype = compute_dtype
+
+    def _use_mlx_full_forward(self, batch):
+        return not self.training and self.mps_model_backend == "mlx_full" and batch.device.type == "mps"
+
+    def mlx_forward_mx(self, raw_audio):
+        from .....bandit_mlx import mlx_forward_bandit_mx
+
+        return mlx_forward_bandit_mx(self, raw_audio, self.mps_model_compute_dtype)
+
+    def forward(self, batch):
+        if self._use_mlx_full_forward(batch):
+            try:
+                from .....bandit_mlx import mlx_forward_bandit
+
+                return mlx_forward_bandit(self, batch, self.mps_model_compute_dtype)
+            except Exception as exc:
+                self._pymss_mlx_full_backend_error = repr(exc)
+                self.mps_model_backend = "torch"
+        with torch.no_grad():
+            x = self.stft(batch)
+
+        length = batch.shape[-1]
+        output = self.bsrnn(x, cond=None)
+        estimates = [self.istft(spec, length) for spec in output["spectrogram"].values()]
+        return torch.stack(estimates, dim=1)
+
+
+class MultiMaskMultiSourceBandSplitRNNSimple(MultiMaskMultiSourceBandSplitBaseSimple):
+    def __init__(
+        self,
+        in_channel: int,
+        stems: List[str],
+        band_specs: Union[str, List[Tuple[float, float]]],
+        fs: int = 44100,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        cond_dim: int = 0,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        n_bands: int = None,
+        use_freq_weights: bool = True,
+        normalize_input: bool = False,
+        mult_add_mask: bool = False,
+        freeze_encoder: bool = False,
+    ) -> None:
+        super().__init__(
+            stems=stems,
+            band_specs=band_specs,
+            fs=fs,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            n_bands=n_bands,
+        )
+
+        self.bsrnn = MultiSourceMultiMaskBandSplitCoreRNN(
+            stems=stems,
+            band_specs=self.band_specs,
+            in_channel=in_channel,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            overlapping_band=self.overlapping_band,
+            freq_weights=self.freq_weights,
+            n_freq=n_fft // 2 + 1,
+            use_freq_weights=use_freq_weights,
+            mult_add_mask=mult_add_mask,
+        )
+
+        self.normalize_input = normalize_input
+        self.cond_dim = cond_dim
+
+        if freeze_encoder:
+            for param in self.bsrnn.band_split.parameters():
+                param.requires_grad = False
+
+            for param in self.bsrnn.tf_model.parameters():
+                param.requires_grad = False

tools/pymss_core/modules/bandit/maskestim.py

@@ -0,0 +1,312 @@
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+from torch import nn
+from torch.nn.modules import activation
+from torch.utils.checkpoint import checkpoint_sequential
+
+from .core.model.bsrnn.utils import (
+    band_widths_from_specs,
+    check_no_gap,
+    check_no_overlap,
+    check_nonzero_bandwidth,
+)
+
+
+def _resolve_channels(in_channels=None, in_channel=None):
+    channels = in_channels if in_channels is not None else in_channel
+    if channels is None:
+        raise TypeError("in_channels is required")
+    return channels
+
+
+class BaseNormMLP(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channels: Optional[int] = None,
+        in_channel: Optional[int] = None,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ):
+        super().__init__()
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+        channels = _resolve_channels(in_channels, in_channel)
+        self.hidden_activation_kwargs = hidden_activation_kwargs
+        self.norm = nn.LayerNorm(emb_dim)
+        self.hidden = nn.Sequential(
+            nn.Linear(in_features=emb_dim, out_features=mlp_dim),
+            activation.__dict__[hidden_activation](**hidden_activation_kwargs),
+        )
+        self.bandwidth = bandwidth
+        self.in_channels = channels
+        self.in_channel = channels
+        self.complex_mask = complex_mask
+        self.reim = 2 if complex_mask else 1
+        self.glu_mult = 2
+
+
+class NormMLP(BaseNormMLP):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channels: Optional[int] = None,
+        in_channel: Optional[int] = None,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+        use_combined: bool = False,
+        use_checkpoint: bool = False,
+    ) -> None:
+        super().__init__(
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            bandwidth=bandwidth,
+            in_channels=in_channels,
+            in_channel=in_channel,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+        self.output = nn.Sequential(
+            nn.Linear(
+                in_features=mlp_dim,
+                out_features=self.bandwidth * self.in_channels * self.reim * 2,
+            ),
+            nn.GLU(dim=-1),
+        )
+        self.use_checkpoint = use_checkpoint
+        if use_combined:
+            self.combined = nn.Sequential(self.norm, self.hidden, self.output)
+
+    def reshape_output(self, mb):
+        batch, n_time, _ = mb.shape
+        if self.complex_mask:
+            mb = torch.view_as_complex(mb.reshape(batch, n_time, self.in_channels, self.bandwidth, self.reim).contiguous())
+        else:
+            mb = mb.reshape(batch, n_time, self.in_channels, self.bandwidth)
+
+        return mb.permute(0, 2, 3, 1)
+
+    def forward(self, qb):
+        if hasattr(self, "combined"):
+            if self.use_checkpoint:
+                mb = checkpoint_sequential(self.combined, 2, qb, use_reentrant=False)
+            else:
+                mb = self.combined(qb)
+        else:
+            mb = self.output(self.hidden(self.norm(qb)))
+        return self.reshape_output(mb)
+
+
+class MultAddNormMLP(NormMLP):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channels: Optional[int] = None,
+        in_channel: Optional[int] = None,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ) -> None:
+        super().__init__(
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            bandwidth=bandwidth,
+            in_channels=in_channels,
+            in_channel=in_channel,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+        self.output2 = nn.Sequential(
+            nn.Linear(
+                in_features=mlp_dim,
+                out_features=self.bandwidth * self.in_channels * self.reim * 2,
+            ),
+            nn.GLU(dim=-1),
+        )
+
+    def forward(self, qb):
+        qb = self.hidden(self.norm(qb))
+        return self.reshape_output(self.output(qb)), self.reshape_output(self.output2(qb))
+
+
+class MaskEstimationModuleSuperBase(nn.Module):
+    pass
+
+
+class MaskEstimationModuleBase(MaskEstimationModuleSuperBase):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        in_channels: Optional[int] = None,
+        in_channel: Optional[int] = None,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        norm_mlp_cls: Type[nn.Module] = NormMLP,
+        norm_mlp_kwargs: Dict = None,
+    ) -> None:
+        super().__init__()
+        channels = _resolve_channels(in_channels, in_channel)
+        self.band_widths = band_widths_from_specs(band_specs)
+        self.n_bands = len(band_specs)
+        hidden_activation_kwargs = hidden_activation_kwargs or {}
+        norm_mlp_kwargs = norm_mlp_kwargs or {}
+        self.norm_mlp = nn.ModuleList(
+            [
+                norm_mlp_cls(
+                    bandwidth=self.band_widths[b],
+                    emb_dim=emb_dim,
+                    mlp_dim=mlp_dim,
+                    in_channels=channels,
+                    hidden_activation=hidden_activation,
+                    hidden_activation_kwargs=hidden_activation_kwargs,
+                    complex_mask=complex_mask,
+                    **norm_mlp_kwargs,
+                )
+                for b in range(self.n_bands)
+            ]
+        )
+
+    def compute_masks(self, q):
+        return [nmlp(q[:, b, :, :]) for b, nmlp in enumerate(self.norm_mlp)]
+
+    def compute_mask(self, q, b):
+        return self.norm_mlp[b](q[:, b, :, :])
+
+
+class OverlappingMaskEstimationModule(MaskEstimationModuleBase):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        freq_weights: List[torch.Tensor],
+        n_freq: int,
+        emb_dim: int,
+        mlp_dim: int,
+        in_channels: Optional[int] = None,
+        in_channel: Optional[int] = None,
+        cond_dim: int = 0,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        norm_mlp_cls: Type[nn.Module] = NormMLP,
+        norm_mlp_kwargs: Dict = None,
+        use_freq_weights: bool = True,
+        register_all_freq_weights: bool = True,
+        allow_cond: bool = True,
+        output_dtype: str = "mask",
+        compute_all_masks: bool = True,
+    ) -> None:
+        check_nonzero_bandwidth(band_specs)
+        check_no_gap(band_specs)
+        if cond_dim > 0 and not allow_cond:
+            raise NotImplementedError
+
+        channels = _resolve_channels(in_channels, in_channel)
+        super().__init__(
+            band_specs=band_specs,
+            emb_dim=emb_dim + cond_dim,
+            mlp_dim=mlp_dim,
+            in_channels=channels,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            norm_mlp_cls=norm_mlp_cls,
+            norm_mlp_kwargs=norm_mlp_kwargs,
+        )
+        self.n_freq = n_freq
+        self.band_specs = band_specs
+        self.in_channels = channels
+        self.in_channel = channels
+        self.cond_dim = cond_dim
+        self.allow_cond = allow_cond
+        self.output_dtype = output_dtype
+        self.compute_all_masks = compute_all_masks
+
+        should_register = freq_weights is not None and (register_all_freq_weights or use_freq_weights)
+        self.use_freq_weights = bool(freq_weights is not None and use_freq_weights)
+        if should_register:
+            for i, fw in enumerate(freq_weights):
+                self.register_buffer(f"freq_weights/{i}", fw)
+
+    def _append_cond(self, q, cond):
+        if cond is not None:
+            batch, n_bands, n_time, _ = q.shape
+            if cond.ndim == 2:
+                cond = cond[:, None, None, :].expand(-1, n_bands, n_time, -1)
+            elif cond.ndim == 3:
+                assert cond.shape[1] == n_time
+            else:
+                raise ValueError(f"Invalid cond shape: {cond.shape}")
+            return torch.cat([q, cond], dim=-1)
+
+        if self.cond_dim <= 0:
+            return q
+
+        batch, n_bands, n_time, _ = q.shape
+        cond = torch.ones(batch, n_bands, n_time, self.cond_dim, device=q.device, dtype=q.dtype)
+        return torch.cat([q, cond], dim=-1)
+
+    def forward(self, q, cond=None):
+        if not self.allow_cond and cond is not None:
+            raise NotImplementedError
+        q = self._append_cond(q, cond)
+        batch, n_bands, n_time, _ = q.shape
+
+        mask_list = self.compute_masks(q) if self.compute_all_masks else None
+        dtype = torch.complex64 if self.output_dtype == "complex64" else mask_list[0].dtype
+        masks = torch.zeros(batch, self.in_channels, self.n_freq, n_time, device=q.device, dtype=dtype)
+
+        for im in range(n_bands):
+            fstart, fend = self.band_specs[im]
+            mask = mask_list[im] if mask_list is not None else self.compute_mask(q, im)
+            if self.use_freq_weights:
+                mask = mask * self.get_buffer(f"freq_weights/{im}")[:, None]
+            masks[:, :, fstart:fend, :] += mask
+
+        return masks
+
+
+class MaskEstimationModule(OverlappingMaskEstimationModule):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        in_channels: Optional[int] = None,
+        in_channel: Optional[int] = None,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        **kwargs,
+    ) -> None:
+        check_nonzero_bandwidth(band_specs)
+        check_no_gap(band_specs)
+        check_no_overlap(band_specs)
+        super().__init__(
+            in_channels=in_channels,
+            in_channel=in_channel,
+            band_specs=band_specs,
+            freq_weights=None,
+            n_freq=0,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+    def forward(self, q, cond=None):
+        return torch.concat(self.compute_masks(q), dim=2)

tools/pymss_core/modules/bandit/tfmodel.py

@@ -0,0 +1,194 @@
+import warnings
+
+import torch
+from torch import nn
+from torch.nn.modules import rnn
+from torch.utils.checkpoint import checkpoint_sequential
+
+
+class TimeFrequencyModellingModule(nn.Module):
+    pass
+
+
+class ResidualRNN(nn.Module):
+    def __init__(
+        self,
+        emb_dim: int,
+        rnn_dim: int,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        use_batch_trick: bool = True,
+        use_layer_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        self.use_layer_norm = use_layer_norm
+        self.norm = (
+            nn.LayerNorm(emb_dim)
+            if use_layer_norm
+            else nn.GroupNorm(
+                num_groups=emb_dim,
+                num_channels=emb_dim,
+            )
+        )
+        self.rnn = rnn.__dict__[rnn_type](
+            input_size=emb_dim,
+            hidden_size=rnn_dim,
+            num_layers=1,
+            batch_first=True,
+            bidirectional=bidirectional,
+        )
+        self.fc = nn.Linear(
+            in_features=rnn_dim * (2 if bidirectional else 1),
+            out_features=emb_dim,
+        )
+        self.use_batch_trick = use_batch_trick
+        if not self.use_batch_trick:
+            warnings.warn("NOT USING BATCH TRICK IS EXTREMELY SLOW!!")
+
+    def forward(self, z):
+        z0 = torch.clone(z)
+
+        if self.use_layer_norm:
+            z = self.norm(z)
+        else:
+            z = self.norm(z.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+
+        batch, n_uncrossed, n_across, emb_dim = z.shape
+
+        if self.use_batch_trick:
+            z = self.rnn(z.reshape(batch * n_uncrossed, n_across, emb_dim).contiguous())[0].reshape(
+                batch, n_uncrossed, n_across, -1
+            )
+        else:
+            z = torch.stack([self.rnn(z[:, i, :, :])[0] for i in range(n_uncrossed)], dim=1)
+
+        return self.fc(z) + z0
+
+
+class Transpose(nn.Module):
+    def __init__(self, dim0: int, dim1: int) -> None:
+        super().__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, z):
+        return z.transpose(self.dim0, self.dim1)
+
+
+class SeqBandModellingModule(TimeFrequencyModellingModule):
+    def __init__(
+        self,
+        n_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        parallel_mode: bool = False,
+        sequential_transpose: bool = False,
+        checkpoint_segments: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.n_modules = n_modules
+        self.parallel_mode = parallel_mode
+        self.checkpoint_segments = checkpoint_segments
+
+        if parallel_mode:
+            self.seqband = nn.ModuleList(
+                [
+                    nn.ModuleList(
+                        [
+                            ResidualRNN(
+                                emb_dim=emb_dim,
+                                rnn_dim=rnn_dim,
+                                bidirectional=bidirectional,
+                                rnn_type=rnn_type,
+                            ),
+                            ResidualRNN(
+                                emb_dim=emb_dim,
+                                rnn_dim=rnn_dim,
+                                bidirectional=bidirectional,
+                                rnn_type=rnn_type,
+                            ),
+                        ]
+                    )
+                    for _ in range(n_modules)
+                ]
+            )
+        elif sequential_transpose:
+            layers = []
+            for _ in range(2 * n_modules):
+                layers.extend(
+                    [
+                        ResidualRNN(
+                            emb_dim=emb_dim,
+                            rnn_dim=rnn_dim,
+                            bidirectional=bidirectional,
+                            rnn_type=rnn_type,
+                        ),
+                        Transpose(1, 2),
+                    ]
+                )
+            self.seqband = nn.Sequential(*layers)
+        else:
+            self.seqband = nn.ModuleList(
+                [
+                    ResidualRNN(
+                        emb_dim=emb_dim,
+                        rnn_dim=rnn_dim,
+                        bidirectional=bidirectional,
+                        rnn_type=rnn_type,
+                    )
+                    for _ in range(2 * n_modules)
+                ]
+            )
+
+    def forward(self, z):
+        if self.parallel_mode:
+            for sbm_t, sbm_f in self.seqband:
+                zt = sbm_t(z)
+                zf = sbm_f(z.transpose(1, 2))
+                z = zt + zf.transpose(1, 2)
+            return z
+
+        if isinstance(self.seqband, nn.Sequential):
+            if self.checkpoint_segments:
+                return checkpoint_sequential(
+                    self.seqband,
+                    self.checkpoint_segments,
+                    z,
+                    use_reentrant=False,
+                )
+            return self.seqband(z)
+
+        for sbm in self.seqband:
+            z = sbm(z)
+            z = z.transpose(1, 2)
+        return z
+
+
+class _SeqBandModellingPreset(SeqBandModellingModule):
+    def __init__(
+        self,
+        n_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        parallel_mode: bool = False,
+    ) -> None:
+        super().__init__(
+            n_modules=n_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            parallel_mode=parallel_mode,
+            **self._preset_runtime_options(n_modules, parallel_mode),
+        )
+
+    @staticmethod
+    def _preset_runtime_options(n_modules, parallel_mode):
+        return {
+            "sequential_transpose": False,
+            "checkpoint_segments": None,
+        }

tools/pymss_core/modules/bandit_mlx.py

@@ -0,0 +1,415 @@
+import numpy as np
+import torch
+
+from .mlx_utils import mlx_periodic_hann_window
+from .bandit.tfmodel import ResidualRNN, Transpose
+from .bs_roformer.mlx_attention import _gelu, _linear, _mlx_dtype, _torch_to_mlx_array, mlx_to_torch_mps
+
+
+def torch_to_mlx_input(tensor, dtype):
+    import mlx.core as mx
+
+    return mx.array(tensor.detach().to(dtype=dtype).cpu().numpy())
+
+
+def _mlx_param(module, name, tensor, dtype):
+    cache = getattr(module, "_pymss_mlx_full_param_cache", None)
+    if cache is None:
+        cache = {}
+        module._pymss_mlx_full_param_cache = cache
+    key = (name, tensor.data_ptr(), tensor._version, tuple(tensor.shape), dtype)
+    cached = cache.get(name)
+    if cached is not None and cached[0] == key:
+        return cached[1]
+    value = _torch_to_mlx_array(tensor, dtype)
+    cache[name] = (key, value)
+    return value
+
+
+def _hann_window(length, dtype):
+    return mlx_periodic_hann_window(length, dtype)
+
+
+def _pad_last(x, left, right, mode="constant"):
+    import mlx.core as mx
+
+    if left <= 0 and right <= 0:
+        return x
+    if mode == "constant":
+        return mx.pad(x, [(0, 0)] * (x.ndim - 1) + [(left, right)])
+    if mode != "reflect":
+        raise TypeError(f"MLX Bandit STFT does not support pad_mode={mode!r}")
+    if x.shape[-1] <= max(left, right):
+        raise ValueError("reflect padding requires input length greater than padding")
+    parts = []
+    if left > 0:
+        parts.append(x[..., 1 : left + 1][..., ::-1])
+    parts.append(x)
+    if right > 0:
+        parts.append(x[..., -right - 1 : -1][..., ::-1])
+    return mx.concatenate(parts, axis=-1)
+
+
+def _spectral_stft(stft_module, raw_audio, dtype):
+    import mlx.core as mx
+
+    leading_shape = raw_audio.shape[:-1]
+    length = raw_audio.shape[-1]
+    n_fft = int(stft_module.n_fft)
+    win_length = int(stft_module.win_length)
+    hop = int(stft_module.hop_length)
+    x = raw_audio.reshape(-1, length).astype(dtype)
+    if stft_module.center:
+        x = _pad_last(x, n_fft // 2, n_fft // 2, stft_module.pad_mode)
+    frames = 1 + (x.shape[-1] - n_fft) // hop
+    framed = mx.as_strided(x, shape=(x.shape[0], frames, n_fft), strides=(x.shape[-1], hop, 1))
+    window = _hann_window(win_length, dtype)
+    if win_length < n_fft:
+        left = (n_fft - win_length) // 2
+        window = mx.pad(window, [(left, n_fft - win_length - left)])
+    elif win_length > n_fft:
+        raise ValueError("MLX Bandit STFT does not support win_length > n_fft")
+    spec = mx.fft.rfft(framed * window, n=n_fft, axis=-1)
+    if stft_module.normalized:
+        spec = spec / np.sqrt(n_fft)
+    spec = mx.moveaxis(spec, -1, -2)
+    return spec.reshape(*leading_shape, spec.shape[-2], spec.shape[-1]), {
+        "n_fft": n_fft,
+        "win_length": win_length,
+        "hop": hop,
+        "window": window,
+        "normalized": stft_module.normalized,
+        "center": stft_module.center,
+        "dtype": dtype,
+    }
+
+
+def _spectral_istft(istft_module, spec, context, length):
+    import mlx.core as mx
+
+    leading_shape = spec.shape[:-2]
+    freqs, frames_n = spec.shape[-2:]
+    n_fft = context["n_fft"]
+    hop = context["hop"]
+    complex_stft = mx.moveaxis(spec.reshape(-1, freqs, frames_n), -2, -1)
+    if context["normalized"]:
+        complex_stft = complex_stft * np.sqrt(n_fft)
+    frames = mx.fft.irfft(complex_stft, n=n_fft, axis=-1).astype(context["dtype"]) * context["window"]
+    full_length = n_fft + hop * (frames.shape[1] - 1)
+    positions = mx.arange(n_fft)[None, :] + hop * mx.arange(frames.shape[1])[:, None]
+    audio = mx.zeros((frames.shape[0], full_length), dtype=context["dtype"]).at[:, positions].add(frames)
+    denom_frames = mx.broadcast_to(mx.square(context["window"])[None, :], (frames.shape[1], n_fft))
+    denom = mx.zeros((full_length,), dtype=context["dtype"]).at[positions].add(denom_frames)
+    audio = audio / mx.maximum(denom[None, :], mx.array(1e-11, dtype=context["dtype"]))
+    if context["center"]:
+        pad = n_fft // 2
+        audio = audio[..., pad : pad + length]
+    elif length is not None:
+        audio = audio[..., :length]
+    return audio.reshape(*leading_shape, audio.shape[-1])
+
+
+def _layer_norm(module, x, dtype):
+    import mlx.core as mx
+
+    x32 = x.astype(mx.float32)
+    mean = mx.mean(x32, axis=-1, keepdims=True)
+    var = mx.mean(mx.square(x32 - mean), axis=-1, keepdims=True)
+    y = ((x32 - mean) * mx.rsqrt(var + module.eps)).astype(x.dtype)
+    if module.elementwise_affine:
+        y = y * _mlx_param(module, "weight", module.weight, dtype)
+        if module.bias is not None:
+            y = y + _mlx_param(module, "bias", module.bias, dtype)
+    return y
+
+
+def _group_norm_nchw(module, x, dtype):
+    import mlx.core as mx
+
+    b, c = x.shape[:2]
+    rest = x.shape[2:]
+    groups = int(module.num_groups)
+    y = x.astype(mx.float32).reshape(b, groups, c // groups, *rest)
+    axes = tuple(range(2, y.ndim))
+    mean = mx.mean(y, axis=axes, keepdims=True)
+    var = mx.mean(mx.square(y - mean), axis=axes, keepdims=True)
+    y = ((y - mean) * mx.rsqrt(var + module.eps)).reshape(x.shape).astype(x.dtype)
+    if module.affine:
+        shape = (1, -1) + (1,) * len(rest)
+        y = y * _mlx_param(module, "weight", module.weight, dtype).reshape(*shape)
+        y = y + _mlx_param(module, "bias", module.bias, dtype).reshape(*shape)
+    return y
+
+
+def _activation(module, x):
+    import mlx.core as mx
+
+    if isinstance(module, torch.nn.Tanh):
+        return mx.tanh(x)
+    if isinstance(module, torch.nn.ReLU):
+        return mx.maximum(x, 0)
+    if isinstance(module, torch.nn.GELU):
+        return _gelu(x)
+    if isinstance(module, torch.nn.ELU):
+        return mx.where(x > 0, x, module.alpha * (mx.exp(x) - 1))
+    if isinstance(module, torch.nn.Identity):
+        return x
+    raise TypeError(f"unsupported Bandit activation for MLX full backend: {type(module).__name__}")
+
+
+def _glu(x, axis=-1):
+    import mlx.core as mx
+
+    a, b = mx.split(x, 2, axis=axis)
+    return a * mx.sigmoid(b)
+
+
+def _norm_fc(module, xb, dtype):
+    if hasattr(module, "combined"):
+        xb = _layer_norm(module.combined[0], xb, dtype)
+        return _linear(
+            xb,
+            _mlx_param(module.combined[1], "weight", module.combined[1].weight, dtype),
+            _mlx_param(module.combined[1], "bias", module.combined[1].bias, dtype),
+        )
+
+    batch, n_time, in_channels, ribw = xb.shape
+    xb = _layer_norm(module.norm, xb.reshape(batch, n_time, in_channels * ribw), dtype)
+    if module.treat_channel_as_feature:
+        return _linear(
+            xb, _mlx_param(module.fc, "weight", module.fc.weight, dtype), _mlx_param(module.fc, "bias", module.fc.bias, dtype)
+        )
+    out = _linear(
+        xb.reshape(batch, n_time, in_channels, ribw),
+        _mlx_param(module.fc, "weight", module.fc.weight, dtype),
+        _mlx_param(module.fc, "bias", module.fc.bias, dtype),
+    )
+    return out.reshape(batch, n_time, -1)
+
+
+def _band_split(module, x, dtype):
+    import mlx.core as mx
+
+    batch, in_channels, _, n_time = x.shape
+    xr = mx.stack((x.real, x.imag), axis=-1)
+    if module.complex_order == "reim_freq":
+        xr = xr.transpose(0, 3, 1, 4, 2)
+    elif module.complex_order == "freq_reim":
+        xr = xr.transpose(0, 3, 1, 2, 4)
+    else:
+        raise ValueError(f"unsupported complex_order: {module.complex_order}")
+
+    outs = []
+    for i, nfm in enumerate(module.norm_fc_modules):
+        fstart, fend = module.band_specs[i]
+        if module.complex_order == "reim_freq":
+            xb = xr[..., fstart:fend].reshape(batch, n_time, in_channels, -1)
+        else:
+            xb = xr[:, :, :, fstart:fend].reshape(batch, n_time, -1)
+        outs.append(_norm_fc(nfm, xb.reshape(batch, n_time, -1) if module.flatten_input else xb, dtype))
+    return mx.stack(outs, axis=1)
+
+
+def _rnn_forward(rnn, x, dtype):
+    import mlx.core as mx
+
+    def params(suffix):
+        return {
+            "w_ih": _mlx_param(rnn, f"weight_ih_l0{suffix}", getattr(rnn, f"weight_ih_l0{suffix}"), dtype),
+            "w_hh": _mlx_param(rnn, f"weight_hh_l0{suffix}", getattr(rnn, f"weight_hh_l0{suffix}"), dtype),
+            "b_ih": _mlx_param(rnn, f"bias_ih_l0{suffix}", getattr(rnn, f"bias_ih_l0{suffix}"), dtype) if rnn.bias else None,
+            "b_hh": _mlx_param(rnn, f"bias_hh_l0{suffix}", getattr(rnn, f"bias_hh_l0{suffix}"), dtype) if rnn.bias else None,
+        }
+
+    def affine(inp, h, p):
+        gates = _linear(inp, p["w_ih"], p["b_ih"]) + _linear(h, p["w_hh"], p["b_hh"])
+        return gates
+
+    def run_gru(inp, p, reverse=False):
+        steps = range(inp.shape[1] - 1, -1, -1) if reverse else range(inp.shape[1])
+        h = mx.zeros((inp.shape[0], rnn.hidden_size), dtype=inp.dtype)
+        outs = []
+        for t in steps:
+            gi = _linear(inp[:, t], p["w_ih"], p["b_ih"])
+            gh = _linear(h, p["w_hh"], p["b_hh"])
+            i_r, i_z, i_n = mx.split(gi, 3, axis=-1)
+            h_r, h_z, h_n = mx.split(gh, 3, axis=-1)
+            reset = mx.sigmoid(i_r + h_r)
+            update = mx.sigmoid(i_z + h_z)
+            new = mx.tanh(i_n + reset * h_n)
+            h = (1 - update) * new + update * h
+            outs.append(h)
+        if reverse:
+            outs.reverse()
+        return mx.stack(outs, axis=1)
+
+    def run_lstm(inp, p, reverse=False):
+        steps = range(inp.shape[1] - 1, -1, -1) if reverse else range(inp.shape[1])
+        h = mx.zeros((inp.shape[0], rnn.hidden_size), dtype=inp.dtype)
+        c = mx.zeros_like(h)
+        outs = []
+        for t in steps:
+            i, f, g, o = mx.split(affine(inp[:, t], h, p), 4, axis=-1)
+            i, f, o = mx.sigmoid(i), mx.sigmoid(f), mx.sigmoid(o)
+            c = f * c + i * mx.tanh(g)
+            h = o * mx.tanh(c)
+            outs.append(h)
+        if reverse:
+            outs.reverse()
+        return mx.stack(outs, axis=1)
+
+    if rnn.num_layers != 1 or not rnn.batch_first:
+        raise TypeError("MLX Bandit RNN supports one-layer batch_first RNNs only")
+    if isinstance(rnn, torch.nn.GRU):
+        forward = run_gru(x, params(""))
+        if not rnn.bidirectional:
+            return forward
+        return mx.concatenate((forward, run_gru(x, params("_reverse"), reverse=True)), axis=-1)
+    if isinstance(rnn, torch.nn.LSTM):
+        forward = run_lstm(x, params(""))
+        if not rnn.bidirectional:
+            return forward
+        return mx.concatenate((forward, run_lstm(x, params("_reverse"), reverse=True)), axis=-1)
+    raise TypeError(f"unsupported Bandit RNN for MLX full backend: {type(rnn).__name__}")
+
+
+def _residual_rnn(module, z, dtype):
+    z0 = z
+    if module.use_layer_norm:
+        z = _layer_norm(module.norm, z, dtype)
+    else:
+        z = _group_norm_nchw(module.norm, z.transpose(0, 3, 1, 2), dtype).transpose(0, 2, 3, 1)
+
+    batch, n_uncrossed, n_across, emb_dim = z.shape
+    if module.use_batch_trick:
+        z = _rnn_forward(module.rnn, z.reshape(batch * n_uncrossed, n_across, emb_dim), dtype)
+        z = z.reshape(batch, n_uncrossed, n_across, -1)
+    else:
+        import mlx.core as mx
+
+        z = mx.stack([_rnn_forward(module.rnn, z[:, i], dtype) for i in range(n_uncrossed)], axis=1)
+    return (
+        _linear(
+            z, _mlx_param(module.fc, "weight", module.fc.weight, dtype), _mlx_param(module.fc, "bias", module.fc.bias, dtype)
+        )
+        + z0
+    )
+
+
+def _tf_model(module, z, dtype):
+    if module.parallel_mode:
+        for sbm_t, sbm_f in module.seqband:
+            zt = _residual_rnn(sbm_t, z, dtype)
+            zf = _residual_rnn(sbm_f, z.transpose(0, 2, 1, 3), dtype)
+            z = zt + zf.transpose(0, 2, 1, 3)
+        return z
+
+    if isinstance(module.seqband, torch.nn.Sequential):
+        for layer in module.seqband:
+            if isinstance(layer, ResidualRNN):
+                z = _residual_rnn(layer, z, dtype)
+            elif isinstance(layer, Transpose):
+                z = z.swapaxes(layer.dim0, layer.dim1)
+            else:
+                raise TypeError(f"unsupported Bandit TF layer for MLX full backend: {type(layer).__name__}")
+        return z
+
+    for sbm in module.seqband:
+        z = _residual_rnn(sbm, z, dtype)
+        z = z.swapaxes(1, 2)
+    return z
+
+
+def _norm_mlp(module, qb, dtype):
+    x = _layer_norm(module.norm, qb, dtype)
+    x = _linear(
+        x,
+        _mlx_param(module.hidden[0], "weight", module.hidden[0].weight, dtype),
+        _mlx_param(module.hidden[0], "bias", module.hidden[0].bias, dtype),
+    )
+    x = _activation(module.hidden[1], x)
+    output = module.output[0]
+    x = _linear(x, _mlx_param(output, "weight", output.weight, dtype), _mlx_param(output, "bias", output.bias, dtype))
+    x = _glu(x, axis=-1)
+
+    batch, n_time, _ = x.shape
+    if module.complex_mask:
+        x = x.reshape(batch, n_time, module.in_channels, module.bandwidth, 2)
+        x = x[..., 0] + (1j * x[..., 1])
+    else:
+        x = x.reshape(batch, n_time, module.in_channels, module.bandwidth)
+    return x.transpose(0, 2, 3, 1)
+
+
+def _append_cond(module, q, cond):
+    import mlx.core as mx
+
+    if cond is not None:
+        batch, n_bands, n_time, _ = q.shape
+        if cond.ndim == 2:
+            cond = mx.broadcast_to(cond[:, None, None, :], (batch, n_bands, n_time, cond.shape[-1]))
+        elif cond.ndim != 3:
+            raise ValueError(f"Invalid cond shape: {cond.shape}")
+        return mx.concatenate((q, cond), axis=-1)
+    if module.cond_dim <= 0:
+        return q
+    batch, n_bands, n_time, _ = q.shape
+    return mx.concatenate((q, mx.ones((batch, n_bands, n_time, module.cond_dim), dtype=q.dtype)), axis=-1)
+
+
+def _mask_estimator(module, q, dtype, cond=None):
+    import mlx.core as mx
+
+    q = _append_cond(module, q, cond)
+    if getattr(module, "n_freq", 0) <= 0:
+        return mx.concatenate([_norm_mlp(nmlp, q[:, b], dtype) for b, nmlp in enumerate(module.norm_mlp)], axis=2)
+
+    batch, _, n_time, _ = q.shape
+    mask_real = mx.zeros((batch, module.in_channels, module.n_freq, n_time), dtype=mx.float32)
+    mask_imag = mx.zeros_like(mask_real)
+    for band_index, nmlp in enumerate(module.norm_mlp):
+        fstart, fend = module.band_specs[band_index]
+        mask = _norm_mlp(nmlp, q[:, band_index], dtype)
+        if module.use_freq_weights:
+            fw = _torch_to_mlx_array(module.get_buffer(f"freq_weights/{band_index}"), dtype)
+            mask = mask * fw.reshape(1, 1, -1, 1)
+        padding = [(0, 0), (0, 0), (fstart, module.n_freq - fend), (0, 0)]
+        mask_real = mask_real + mx.pad(mask.real.astype(mask_real.dtype), padding)
+        mask_imag = mask_imag + mx.pad(mask.imag.astype(mask_imag.dtype), padding)
+    return mask_real + (1j * mask_imag)
+
+
+def _bsrnn_core(module, x, dtype):
+    batch, in_chan, n_freq, n_time = x.shape
+    x = x.reshape(-1, 1, n_freq, n_time)
+    q = _tf_model(module.tf_model, _band_split(module.band_split, x, dtype), dtype)
+    return [_mask_estimator(mask_estimator, q, dtype) * x for mask_estimator in module.mask_estim.values()]
+
+
+def mlx_forward_bandit_mx(module, raw_audio, dtype=torch.float16):
+    import mlx.core as mx
+
+    if dtype not in (torch.float16, torch.float32):
+        raise TypeError("MLX full Bandit supports torch.float16 or torch.float32 compute dtype")
+    mx_dtype = _mlx_dtype(dtype)
+    init_shape = raw_audio.shape
+    mono = raw_audio.reshape(-1, 1, raw_audio.shape[-1]).astype(mx_dtype)
+    x, context = _spectral_stft(module.stft, mono, mx_dtype)
+    length = mono.shape[-1]
+
+    if hasattr(module, "bsrnn"):
+        specs = _bsrnn_core(module.bsrnn, x, dtype)
+        stems = module.stems
+    else:
+        q = _tf_model(module.tf_model, _band_split(module.band_split, x, dtype), dtype)
+        specs = [_mask_estimator(mask_estimator, q, dtype) * x for mask_estimator in module.mask_estim.values()]
+        stems = module.stems
+
+    estimates = [_spectral_istft(module.istft, spec, context, length) for spec in specs]
+    estimates = [estimate.reshape(-1, init_shape[1], init_shape[2]) for estimate in estimates]
+    return mx.stack(estimates, axis=1)
+
+
+def mlx_forward_bandit(module, raw_audio, dtype=torch.float16):
+    x_mx = torch_to_mlx_input(raw_audio, dtype=dtype)
+    return mlx_to_torch_mps(mlx_forward_bandit_mx(module, x_mx, dtype), raw_audio)

tools/pymss_core/modules/bandit_v2/bandit.py

@@ -0,0 +1,326 @@
+from typing import Dict, List, Optional
+
+import torch
+from torch import nn
+
+from ..bandit.core.model._spectral import _SpectralComponent
+
+from .bandsplit import BandSplitModule
+from .maskestim import OverlappingMaskEstimationModule
+from .tfmodel import SeqBandModellingModule
+from .utils import MusicalBandsplitSpecification
+
+
+class BaseBandit(_SpectralComponent):
+    mps_model_backend = "torch"
+    mps_model_compute_dtype = torch.float16
+
+    def __init__(
+        self,
+        in_channels: int,
+        fs: int,
+        band_type: str = "musical",
+        n_bands: int = 64,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        n_fft: int = 2048,
+        win_length: Optional[int] = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Optional[Dict] = None,
+        power: Optional[int] = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+    ):
+        super().__init__(
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            normalized=normalized,
+            center=center,
+            pad_mode=pad_mode,
+            onesided=onesided,
+        )
+        self.in_channels = in_channels
+
+        self.instantiate_bandsplit(
+            in_channels=in_channels,
+            band_type=band_type,
+            n_bands=n_bands,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+            n_fft=n_fft,
+            fs=fs,
+        )
+
+        self.instantiate_tf_modelling(
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+        )
+
+    def set_mps_model_backend(self, backend=None, compute_dtype=None):
+        backend = (backend or "torch").lower()
+        if backend not in ("torch", "mlx_full"):
+            raise ValueError("mps_model_backend must be 'torch' or 'mlx_full'")
+        self.mps_model_backend = backend
+        if compute_dtype is None:
+            return
+        if isinstance(compute_dtype, str):
+            compute_dtype = {
+                "float16": torch.float16,
+                "fp16": torch.float16,
+                "float32": torch.float32,
+                "fp32": torch.float32,
+            }.get(compute_dtype.lower(), compute_dtype)
+        if compute_dtype not in (torch.float16, torch.float32):
+            raise ValueError("mps_model_compute_dtype must be 'float16' or 'float32'")
+        self.mps_model_compute_dtype = compute_dtype
+
+    def _use_mlx_full_forward(self, batch):
+        return (
+            not self.training
+            and self.mps_model_backend == "mlx_full"
+            and not isinstance(batch, dict)
+            and batch.device.type == "mps"
+        )
+
+    def mlx_forward_mx(self, raw_audio):
+        from ..bandit_mlx import mlx_forward_bandit_mx
+
+        return mlx_forward_bandit_mx(self, raw_audio, self.mps_model_compute_dtype)
+
+    def instantiate_bandsplit(
+        self,
+        in_channels: int,
+        band_type: str = "musical",
+        n_bands: int = 64,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        emb_dim: int = 128,
+        n_fft: int = 2048,
+        fs: int = 44100,
+    ):
+        assert band_type == "musical"
+
+        self.band_specs = MusicalBandsplitSpecification(nfft=n_fft, fs=fs, n_bands=n_bands)
+
+        self.band_split = BandSplitModule(
+            in_channels=in_channels,
+            band_specs=self.band_specs.get_band_specs(),
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            emb_dim=emb_dim,
+        )
+
+    def instantiate_tf_modelling(
+        self,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+    ):
+        self.tf_model = SeqBandModellingModule(
+            n_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+        )
+
+    def mask(self, x, m):
+        return x * m
+
+    def forward(self, batch, mode="train"):
+        if self._use_mlx_full_forward(batch):
+            try:
+                from ..bandit_mlx import mlx_forward_bandit
+
+                return mlx_forward_bandit(self, batch, self.mps_model_compute_dtype)
+            except Exception as exc:
+                self._pymss_mlx_full_backend_error = repr(exc)
+                self.mps_model_backend = "torch"
+        init_shape = batch.shape
+        if not isinstance(batch, dict):
+            mono = batch.view(-1, 1, batch.shape[-1])
+            batch = {"mixture": {"audio": mono}}
+
+        with torch.no_grad():
+            mixture = batch["mixture"]["audio"]
+
+            x = self.stft(mixture)
+            batch["mixture"]["spectrogram"] = x
+
+            if "sources" in batch.keys():
+                for stem in batch["sources"].keys():
+                    s = batch["sources"][stem]["audio"]
+                    s = self.stft(s)
+                    batch["sources"][stem]["spectrogram"] = s
+
+        batch = self.separate(batch)
+
+        return torch.stack([batch["estimates"][s]["audio"].view(-1, init_shape[1], init_shape[2]) for s in self.stems], dim=1)
+
+    def encode(self, batch):
+        x = batch["mixture"]["spectrogram"]
+        length = batch["mixture"]["audio"].shape[-1]
+
+        z = self.band_split(x)  # (batch, emb_dim, n_band, n_time)
+        q = self.tf_model(z)  # (batch, emb_dim, n_band, n_time)
+
+        return x, q, length
+
+    def separate(self, batch):
+        raise NotImplementedError
+
+
+class Bandit(BaseBandit):
+    def __init__(
+        self,
+        in_channels: int,
+        stems: List[str],
+        band_type: str = "musical",
+        n_bands: int = 64,
+        require_no_overlap: bool = False,
+        require_no_gap: bool = True,
+        normalize_channel_independently: bool = False,
+        treat_channel_as_feature: bool = True,
+        n_sqm_modules: int = 12,
+        emb_dim: int = 128,
+        rnn_dim: int = 256,
+        bidirectional: bool = True,
+        rnn_type: str = "LSTM",
+        mlp_dim: int = 512,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict | None = None,
+        complex_mask: bool = True,
+        use_freq_weights: bool = True,
+        n_fft: int = 2048,
+        win_length: int | None = 2048,
+        hop_length: int = 512,
+        window_fn: str = "hann_window",
+        wkwargs: Dict | None = None,
+        power: int | None = None,
+        center: bool = True,
+        normalized: bool = True,
+        pad_mode: str = "constant",
+        onesided: bool = True,
+        fs: int = 44100,
+        stft_precisions="32",
+        bandsplit_precisions="bf16",
+        tf_model_precisions="bf16",
+        mask_estim_precisions="bf16",
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            band_type=band_type,
+            n_bands=n_bands,
+            require_no_overlap=require_no_overlap,
+            require_no_gap=require_no_gap,
+            normalize_channel_independently=normalize_channel_independently,
+            treat_channel_as_feature=treat_channel_as_feature,
+            n_sqm_modules=n_sqm_modules,
+            emb_dim=emb_dim,
+            rnn_dim=rnn_dim,
+            bidirectional=bidirectional,
+            rnn_type=rnn_type,
+            n_fft=n_fft,
+            win_length=win_length,
+            hop_length=hop_length,
+            window_fn=window_fn,
+            wkwargs=wkwargs,
+            power=power,
+            center=center,
+            normalized=normalized,
+            pad_mode=pad_mode,
+            onesided=onesided,
+            fs=fs,
+        )
+
+        self.stems = stems
+
+        self.instantiate_mask_estim(
+            in_channels=in_channels,
+            stems=stems,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            n_freq=n_fft // 2 + 1,
+            use_freq_weights=use_freq_weights,
+        )
+
+    def instantiate_mask_estim(
+        self,
+        in_channels: int,
+        stems: List[str],
+        emb_dim: int,
+        mlp_dim: int,
+        hidden_activation: str,
+        hidden_activation_kwargs: Optional[Dict] = None,
+        complex_mask: bool = True,
+        n_freq: Optional[int] = None,
+        use_freq_weights: bool = False,
+    ):
+        if hidden_activation_kwargs is None:
+            hidden_activation_kwargs = {}
+
+        assert n_freq is not None
+
+        self.mask_estim = nn.ModuleDict(
+            {
+                stem: OverlappingMaskEstimationModule(
+                    band_specs=self.band_specs.get_band_specs(),
+                    freq_weights=self.band_specs.get_freq_weights(),
+                    n_freq=n_freq,
+                    emb_dim=emb_dim,
+                    mlp_dim=mlp_dim,
+                    in_channels=in_channels,
+                    hidden_activation=hidden_activation,
+                    hidden_activation_kwargs=hidden_activation_kwargs,
+                    complex_mask=complex_mask,
+                    use_freq_weights=use_freq_weights,
+                )
+                for stem in stems
+            }
+        )
+
+    def separate(self, batch):
+        batch["estimates"] = {}
+
+        x, q, length = self.encode(batch)
+
+        for stem, mem in self.mask_estim.items():
+            m = mem(q)
+
+            s = self.mask(x, m.to(x.dtype))
+            s = s.reshape(x.shape)
+            batch["estimates"][stem] = {
+                "audio": self.istft(s, length),
+                "spectrogram": s,
+            }
+
+        return batch

tools/pymss_core/modules/bandit_v2/bandsplit.py

@@ -0,0 +1,13 @@
+from ..bandit.bandsplit import (
+    SequentialNormFC as NormFC,
+    _ConfiguredBandSplitModule,
+)
+
+
+class BandSplitModule(_ConfiguredBandSplitModule):
+    norm_fc_cls = NormFC
+    complex_order = "freq_reim"
+    flatten_input = True
+
+
+__all__ = ("BandSplitModule", "NormFC")

tools/pymss_core/modules/bandit_v2/maskestim.py

@@ -0,0 +1,119 @@
+from typing import Dict, List, Optional, Tuple, Type
+
+import torch
+from torch import nn
+
+from ..bandit.maskestim import (
+    BaseNormMLP,
+    MaskEstimationModule as _MaskEstimationModule,
+    MaskEstimationModuleBase,
+    MaskEstimationModuleSuperBase,
+    NormMLP as _NormMLP,
+    OverlappingMaskEstimationModule as _OverlappingMaskEstimationModule,
+)
+from ..bandit.core.model.bsrnn.utils import (
+    check_no_gap,
+    check_no_overlap,
+    check_nonzero_bandwidth,
+)
+
+
+class NormMLP(_NormMLP):
+    def __init__(
+        self,
+        emb_dim: int,
+        mlp_dim: int,
+        bandwidth: int,
+        in_channels: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs=None,
+        complex_mask: bool = True,
+    ) -> None:
+        super().__init__(
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            bandwidth=bandwidth,
+            in_channels=in_channels,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            use_combined=True,
+            use_checkpoint=True,
+        )
+
+
+class OverlappingMaskEstimationModule(_OverlappingMaskEstimationModule):
+    def __init__(
+        self,
+        in_channels: int,
+        band_specs: List[Tuple[float, float]],
+        freq_weights: List[torch.Tensor],
+        n_freq: int,
+        emb_dim: int,
+        mlp_dim: int,
+        cond_dim: int = 0,
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        norm_mlp_cls: Type[nn.Module] = NormMLP,
+        norm_mlp_kwargs: Dict = None,
+        use_freq_weights: bool = False,
+    ) -> None:
+        super().__init__(
+            in_channels=in_channels,
+            band_specs=band_specs,
+            freq_weights=freq_weights,
+            n_freq=n_freq,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            cond_dim=cond_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+            norm_mlp_cls=norm_mlp_cls,
+            norm_mlp_kwargs=norm_mlp_kwargs,
+            use_freq_weights=use_freq_weights,
+            register_all_freq_weights=False,
+            allow_cond=False,
+            output_dtype="complex64",
+            compute_all_masks=False,
+        )
+
+    def forward(self, q):
+        return super().forward(q)
+
+
+class MaskEstimationModule(_MaskEstimationModule):
+    def __init__(
+        self,
+        band_specs: List[Tuple[float, float]],
+        emb_dim: int,
+        mlp_dim: int,
+        in_channels: Optional[int],
+        hidden_activation: str = "Tanh",
+        hidden_activation_kwargs: Dict = None,
+        complex_mask: bool = True,
+        **kwargs,
+    ) -> None:
+        check_nonzero_bandwidth(band_specs)
+        check_no_gap(band_specs)
+        check_no_overlap(band_specs)
+        super().__init__(
+            in_channels=in_channels,
+            band_specs=band_specs,
+            emb_dim=emb_dim,
+            mlp_dim=mlp_dim,
+            hidden_activation=hidden_activation,
+            hidden_activation_kwargs=hidden_activation_kwargs,
+            complex_mask=complex_mask,
+        )
+
+
+__all__ = (
+    "BaseNormMLP",
+    "MaskEstimationModule",
+    "MaskEstimationModuleBase",
+    "MaskEstimationModuleSuperBase",
+    "NormMLP",
+    "OverlappingMaskEstimationModule",
+)

tools/pymss_core/modules/bandit_v2/tfmodel.py

@@ -0,0 +1,18 @@
+from ..bandit.tfmodel import (
+    ResidualRNN,
+    TimeFrequencyModellingModule,
+    Transpose,
+    _SeqBandModellingPreset,
+)
+
+
+class SeqBandModellingModule(_SeqBandModellingPreset):
+    @staticmethod
+    def _preset_runtime_options(n_modules, parallel_mode):
+        return {
+            "sequential_transpose": not parallel_mode,
+            "checkpoint_segments": None if parallel_mode else n_modules,
+        }
+
+
+__all__ = ("ResidualRNN", "SeqBandModellingModule", "TimeFrequencyModellingModule", "Transpose")