Retrieval-based-Voice-Conversion-WebUI/tools/pymss_webui.py

import gc
import json
import logging
import os
import subprocess
import sys
import tempfile
import threading
import time
import traceback
import uuid
from collections import deque
from concurrent.futures import ThreadPoolExecutor, wait
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

from configs.config import Config
from tools.process_utils import kill_process_tree


logger = logging.getLogger(__name__)
config = Config()
weight_pymss_root = Path(os.getenv("weight_pymss_root", "assets/pymss_weights"))
tools_root = str(Path(__file__).resolve().parent)
if tools_root not in sys.path:
    sys.path.insert(0, tools_root)

MODEL_SAMPLE_RATE = 44100
DML_CHUNK_SIZE = 88200
DML_OVERLAP_SIZE = 22050
DML_MODEL_DTYPE_ENV = "PYMSS_DML_MODEL_DTYPE"
PROJECT_ROOT = Path(__file__).resolve().parents[1]
FFMPEG_PATH = PROJECT_ROOT / "ffmpeg.exe"
AUDIO_PARAMS = {
    "wav_bit_depth": "FLOAT",
    "flac_bit_depth": "PCM_24",
    "mp3_bit_rate": "320k",
    "m4a_bit_rate": "320k",
    "m4a_codec": "aac",
    "m4a_aac_at_quality": 2,
}


@dataclass(frozen=True)
class ModelSpec:
    label: str
    model_id: str
    model_type: str
    model_file: str
    config_file: str
    desired_stem: str
    secondary_stem: str
    desired_suffix: str
    secondary_suffix: str
    batch_size: int
    overlap_size: int


MODEL_SPECS = (
    ModelSpec(
        label="去混响",
        model_id="dereverb-less-aggressive-18.8050",
        model_type="mel_band_roformer",
        model_file="dereverb_mel_band_roformer_less_aggressive_anvuew_sdr_18.8050.ckpt",
        config_file="dereverb_mel_band_roformer_anvuew.yaml",
        desired_stem="noreverb",
        secondary_stem="reverb",
        desired_suffix="noreverb",
        secondary_suffix="reverb",
        batch_size=1,
        overlap_size=176400,
    ),
    ModelSpec(
        label="去混响（激进）",
        model_id="dereverb-anvuew-19.1729",
        model_type="mel_band_roformer",
        model_file="dereverb_mel_band_roformer_anvuew_sdr_19.1729.ckpt",
        config_file="dereverb_mel_band_roformer_anvuew.yaml",
        desired_stem="noreverb",
        secondary_stem="reverb",
        desired_suffix="noreverb",
        secondary_suffix="reverb",
        batch_size=1,
        overlap_size=176400,
    ),
    ModelSpec(
        label="去伴奏",
        model_id="vocals-bs-roformer-368",
        model_type="bs_roformer",
        model_file="model_bs_roformer_ep_368_sdr_12.9628.ckpt",
        config_file="model_bs_roformer_ep_368_sdr_12.9628.yaml",
        desired_stem="vocals",
        secondary_stem="instrumental",
        desired_suffix="vocals",
        secondary_suffix="instrumental",
        batch_size=1,
        overlap_size=264600,
    ),
    ModelSpec(
        label="去伴奏（激进）",
        model_id="vocals-bs-roformer-317",
        model_type="bs_roformer",
        model_file="model_bs_roformer_ep_317_sdr_12.9755.ckpt",
        config_file="model_bs_roformer_ep_317_sdr_12.9755.yaml",
        desired_stem="vocals",
        secondary_stem="other",
        desired_suffix="vocals",
        secondary_suffix="instrumental",
        batch_size=4,
        overlap_size=176400,
    ),
    ModelSpec(
        label="提主旋律",
        model_id="karaoke-mel-roformer-10.1956",
        model_type="mel_band_roformer",
        model_file="model_mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt",
        config_file="config_mel_band_roformer_karaoke.yaml",
        desired_stem="karaoke",
        secondary_stem="other",
        desired_suffix="main_vocal",
        secondary_suffix="off_vocal",
        batch_size=1,
        overlap_size=264600,
    ),
)

MODEL_BY_LABEL = {spec.label: spec for spec in MODEL_SPECS}
MODEL_BY_ID = {spec.model_id: spec for spec in MODEL_SPECS}
PYMSS_MODEL_CHOICES = [spec.label for spec in MODEL_SPECS]
PYMSS_WORKER_STATE_LOCK = threading.Lock()
PYMSS_WORKER_STATE = {
    "task_id": 0,
    "active": False,
    "process": None,
    "stop_requested": False,
}
PYMSS_WORKER_OUTPUT_LOCK = threading.Lock()
DML_FP16_DISABLED_MODEL_TYPES = set()


class DMLFP16Fallback(RuntimeError):
    """Request one retry in a fresh DirectML FP32 worker."""


def _normalize_dml_model_dtype(value):
    normalized = str(value or "auto").strip().lower().replace("torch.", "")
    aliases = {
        "auto": "auto",
        "fp16": "float16",
        "half": "float16",
        "float16": "float16",
        "fp32": "float32",
        "float": "float32",
        "float32": "float32",
    }
    if normalized not in aliases:
        logger.warning(
            "Invalid %s=%r; using auto",
            DML_MODEL_DTYPE_ENV,
            value,
        )
        return "auto"
    return aliases[normalized]


def _dml_model_dtype_attempts(spec):
    requested = _normalize_dml_model_dtype(
        os.getenv(DML_MODEL_DTYPE_ENV, "auto")
    )
    if requested == "float32":
        return ("float32",)
    if requested == "float16":
        return ("float16",)
    if spec.model_type in DML_FP16_DISABLED_MODEL_TYPES:
        return ("float32",)
    return ("float16", "float32")


def resolve_model(model_name):
    if not model_name:
        return MODEL_SPECS[0]
    spec = MODEL_BY_LABEL.get(model_name) or MODEL_BY_ID.get(model_name)
    if spec is None:
        raise ValueError("Unknown separation model: %s" % model_name)
    return spec


def get_model_info(model_name):
    spec = resolve_model(model_name)
    return "%s | %s" % (spec.model_type, spec.model_id)


def clean_path(path):
    path = path or ""
    if path.endswith(("\\", "/")):
        path = path[:-1]
    return path.replace("/", os.sep).replace("\\", os.sep).strip(" '\n\"\u202a")


def _uploaded_path(item):
    if isinstance(item, (str, os.PathLike)):
        return os.fspath(item)
    if isinstance(item, dict):
        return item.get("name") or item.get("path")
    return getattr(item, "name", None)


def collect_input_paths(inp_root, paths):
    inp_root = clean_path(inp_root)
    if inp_root:
        if os.path.isfile(inp_root):
            candidates = [inp_root]
        elif os.path.isdir(inp_root):
            candidates = [os.path.join(inp_root, name) for name in sorted(os.listdir(inp_root))]
        else:
            raise FileNotFoundError(inp_root)
    else:
        candidates = [_uploaded_path(item) for item in (paths or [])]
    return [os.path.abspath(path) for path in candidates if path and os.path.isfile(path)]


def _write_audio(path, audio, sample_rate, output_format):
    audio = np.ascontiguousarray(audio, dtype=np.float32)
    if audio.ndim == 1:
        channels = 1
    elif audio.ndim == 2 and audio.shape[1] in (1, 2):
        channels = audio.shape[1]
    else:
        raise ValueError("Unsupported audio shape: %s" % (audio.shape,))

    if output_format == "wav":
        sf.write(path, audio, sample_rate, format="WAV", subtype="FLOAT")
        return
    if output_format == "flac":
        sf.write(path, audio, sample_rate, format="FLAC", subtype="PCM_24")
        return

    ffmpeg = str(FFMPEG_PATH) if FFMPEG_PATH.is_file() else "ffmpeg"
    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-y",
        "-f",
        "f32le",
        "-ar",
        str(sample_rate),
        "-ac",
        str(channels),
        "-i",
        "pipe:0",
        "-vn",
    ]
    if output_format == "mp3":
        command.extend(("-c:a", "libmp3lame", "-b:a", "320k"))
    elif output_format == "m4a":
        command.extend(("-c:a", "aac", "-aac_coder", "fast", "-b:a", "320k"))
    else:
        raise ValueError("Unsupported output format: %s" % output_format)
    command.append(path)
    completed = subprocess.run(
        command,
        input=audio.tobytes(),
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
    )
    if completed.returncode != 0:
        detail = completed.stderr.decode("utf-8", errors="replace").strip()
        raise RuntimeError("FFmpeg audio encoding failed: %s" % detail)


class MSSTBatchSeparator:
    def __init__(
        self,
        spec,
        output_format,
        desired_root,
        secondary_root,
        model_dtype="auto",
        progress_callback=None,
        separation_logger=None,
    ):
        try:
            from pymss import MSSeparator, load_audio
        except ImportError as error:
            raise RuntimeError(
                "项目内置的 pymss 运行库加载失败"
            ) from error

        self.spec = spec
        self.output_format = output_format.lower()
        if self.output_format not in {"wav", "flac", "mp3", "m4a"}:
            raise ValueError("Unsupported output format: %s" % output_format)
        desired_root = clean_path(desired_root)
        secondary_root = clean_path(secondary_root)
        if not desired_root or not secondary_root:
            raise ValueError("输出文件夹不能为空")
        self.desired_root = os.path.abspath(desired_root)
        self.secondary_root = os.path.abspath(secondary_root)
        os.makedirs(self.desired_root, exist_ok=True)
        os.makedirs(self.secondary_root, exist_ok=True)

        model_path = weight_pymss_root / spec.model_file
        config_path = weight_pymss_root / spec.config_file
        if not model_path.is_file():
            raise FileNotFoundError(model_path)
        if not config_path.is_file():
            raise FileNotFoundError(config_path)

        parsed_device = torch.device(config.device)
        use_cuda = parsed_device.type == "cuda"
        use_dml = parsed_device.type == "privateuseone"
        model_dtype = _normalize_dml_model_dtype(model_dtype) if use_dml else "auto"
        device_id = parsed_device.index if use_cuda and parsed_device.index is not None else 0
        pymss_device = str(parsed_device) if use_dml else ("cuda" if use_cuda else "cpu")
        batch_size = 1 if use_dml else spec.batch_size
        chunk_size = DML_CHUNK_SIZE if use_dml else 352800
        overlap_size = DML_OVERLAP_SIZE if use_dml else spec.overlap_size
        self._load_audio = load_audio
        self.model_load_count = 0
        self.separator = None
        self._save_pool = None
        try:
            self.separator = MSSeparator(
                model_type=spec.model_type,
                model_path=str(model_path),
                config_path=str(config_path),
                device=pymss_device,
                device_ids=[device_id],
                output_format=self.output_format,
                use_tta=False,
                store_dirs={},
                audio_params=AUDIO_PARAMS,
                logger=separation_logger,
                debug=False,
                progress_callback=progress_callback,
                inference_params={
                    "batch_size": batch_size,
                    "chunk_size": chunk_size,
                    "overlap_size": overlap_size,
                    "standardize": False,
                    "normalize": False,
                    "use_amp": bool(config.is_half and use_cuda),
                    "model_dtype": model_dtype,
                    "cuda_attention_backend": "default",
                },
            )
            self.separator.config.training.use_amp = bool(config.is_half and use_cuda)
            model = getattr(self.separator, "model", None)
            model = getattr(model, "module", model)
            parameter = next(model.parameters(), None) if model is not None else None
            self.model_dtype = (
                str(parameter.dtype).replace("torch.", "")
                if parameter is not None
                else "unknown"
            )
            if use_dml and model_dtype == "float16" and self.model_dtype != "float16":
                raise RuntimeError(
                    "DirectML FP16 was requested, but the loaded model uses %s"
                    % self.model_dtype
                )
            self._save_pool = ThreadPoolExecutor(
                max_workers=2,
                thread_name_prefix="rvc-msst-save",
            )
        except Exception:
            self.close()
            raise
        self.model_load_count = 1
        logger.info(
            "Loaded MSST model once for batch: %s, device=%s, model_dtype=%s, amp=%s",
            spec.model_id,
            self.separator.device,
            self.model_dtype,
            bool(config.is_half and use_cuda),
        )

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback_value):
        self.close()

    def _save_output(self, audio, sample_rate, output_root, file_stem, suffix):
        output_path = os.path.join(
            output_root,
            "%s_%s.%s" % (file_stem, suffix, self.output_format),
        )
        temp_path = os.path.join(
            output_root,
            ".%s_%s.%s.tmp.%s"
            % (file_stem, suffix, uuid.uuid4().hex, self.output_format),
        )
        started = time.perf_counter()
        try:
            _write_audio(temp_path, audio, sample_rate, self.output_format)
            if not os.path.isfile(temp_path) or os.path.getsize(temp_path) == 0:
                raise RuntimeError("音频编码没有生成有效文件: %s" % output_path)
            os.replace(temp_path, output_path)
        except Exception:
            if os.path.exists(temp_path):
                os.remove(temp_path)
            raise
        return output_path, time.perf_counter() - started

    def separate_file(self, input_path):
        mix, sample_rate = self._load_audio(input_path, sr=MODEL_SAMPLE_RATE, mono=False)
        inference_started = time.perf_counter()
        results = self.separator.separate(mix, pbar=False)
        inference_seconds = time.perf_counter() - inference_started
        missing = {
            self.spec.desired_stem,
            self.spec.secondary_stem,
        }.difference(results)
        if missing:
            raise RuntimeError("模型缺少输出 stem: %s" % ", ".join(sorted(missing)))
        for stem in (self.spec.desired_stem, self.spec.secondary_stem):
            if not np.isfinite(np.asarray(results[stem])).all():
                raise FloatingPointError("模型输出包含 NaN/Inf: %s" % stem)

        file_stem = Path(input_path).stem
        encode_started = time.perf_counter()
        futures = (
            self._save_pool.submit(
                self._save_output,
                results[self.spec.desired_stem],
                sample_rate,
                self.desired_root,
                file_stem,
                self.spec.desired_suffix,
            ),
            self._save_pool.submit(
                self._save_output,
                results[self.spec.secondary_stem],
                sample_rate,
                self.secondary_root,
                file_stem,
                self.spec.secondary_suffix,
            ),
        )
        wait(futures)
        outputs = [future.result() for future in futures]
        encode_seconds = time.perf_counter() - encode_started
        del results, mix
        return {
            "outputs": [path for path, _ in outputs],
            "inference_seconds": inference_seconds,
            "encode_seconds": encode_seconds,
        }

    def close(self):
        save_pool, self._save_pool = getattr(self, "_save_pool", None), None
        separator, self.separator = getattr(self, "separator", None), None
        try:
            if save_pool is not None:
                try:
                    save_pool.shutdown(wait=True)
                except Exception:
                    logger.exception("Failed to shut down MSST output workers")
            if separator is not None:
                try:
                    separator.close()
                except Exception:
                    logger.exception("Failed to close MSST separator")
        finally:
            del save_pool, separator
            try:
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except Exception:
                logger.exception("Failed to clear Torch cache after MSST task")


def _worker_emit(event):
    with PYMSS_WORKER_OUTPUT_LOCK:
        sys.stdout.write(json.dumps(event, ensure_ascii=False) + "\n")
        sys.stdout.flush()


class _WorkerEventLogHandler(logging.Handler):
    def emit(self, record):
        try:
            _worker_emit(
                {
                    "event": "log",
                    "level": record.levelname,
                    "message": self.format(record),
                }
            )
        except Exception:
            self.handleError(record)


def _looks_like_dml_oom(error):
    message = str(error).lower()
    return any(
        marker in message
        for marker in (
            "out of memory",
            "e_outofmemory",
            "not enough memory",
            "failed to allocate",
            "allocation failed",
        )
    )


def _fp16_retryable(error, requested_dtype, allow_fp32_retry):
    return (
        requested_dtype == "float16"
        and allow_fp32_retry
        and isinstance(
            error,
            (RuntimeError, NotImplementedError, FloatingPointError, TypeError, ValueError),
        )
        and not _looks_like_dml_oom(error)
    )


def _pymss_worker_main(request_path):
    request = {}
    use_dml = False
    requested_dtype = "auto"
    allow_fp32_retry = False
    try:
        with open(request_path, "r", encoding="utf-8") as request_file:
            request = json.load(request_file)
        device_type = torch.device(config.device).type
        use_dml = device_type == "privateuseone"

        spec = resolve_model(request["model_id"])
        input_paths = [os.path.abspath(path) for path in request["input_paths"]]
        if not input_paths:
            raise ValueError("没有找到可处理的音频文件")
        requested_dtype = (
            _normalize_dml_model_dtype(request.get("model_dtype", "float32"))
            if use_dml
            else "auto"
        )
        allow_fp32_retry = use_dml and bool(
            request.get("allow_fp32_retry", False)
        )
        successful_files = 0
        failed_files = 0
        file_count = len(input_paths)
        progress_context = {"file_index": 0, "path": ""}
        progress_emit_state = {"time": 0.0, "message": None, "total": None}

        def emit_progress(done, total, message):
            done = max(0.0, float(done or 0))
            total = max(1.0, float(total or 1))
            done = min(done, total)
            message = str(message or "正在处理音频")
            now = time.monotonic()
            is_edge = done <= 0 or done >= total
            context_changed = (
                progress_emit_state["message"] != message
                or progress_emit_state["total"] != total
            )
            if (
                not is_edge
                and not context_changed
                and now - progress_emit_state["time"] < 0.1
            ):
                return
            progress_emit_state.update(
                {"time": now, "message": message, "total": total}
            )
            _worker_emit(
                {
                    "event": "progress",
                    "file_index": progress_context["file_index"],
                    "file_count": file_count,
                    "path": progress_context["path"],
                    "done": done,
                    "total": total,
                    "message": message,
                }
            )

        worker_logger = logging.getLogger("rvc.pymss.worker")
        worker_logger.handlers.clear()
        worker_logger.setLevel(logging.INFO)
        worker_logger.propagate = False
        worker_log_handler = _WorkerEventLogHandler()
        worker_log_handler.setLevel(logging.INFO)
        worker_log_handler.setFormatter(logging.Formatter("%(message)s"))
        worker_logger.addHandler(worker_log_handler)

        with MSSTBatchSeparator(
            spec,
            request["output_format"],
            request["desired_root"],
            request["secondary_root"],
            model_dtype=requested_dtype,
            progress_callback=emit_progress,
            separation_logger=worker_logger,
        ) as batch:
            device_label = "DirectML" if use_dml else device_type.upper()
            _worker_emit(
                {
                    "event": "status",
                    "message": "%s 模型已加载 | 参数精度 %s"
                    % (device_label, batch.model_dtype.upper()),
                    "model_dtype": batch.model_dtype,
                }
            )
            for file_index, input_path in enumerate(input_paths, 1):
                progress_context.update(
                    {"file_index": file_index, "path": input_path}
                )
                progress_emit_state.update(
                    {"time": 0.0, "message": None, "total": None}
                )
                _worker_emit(
                    {
                        "event": "file_start",
                        "file_index": file_index,
                        "file_count": file_count,
                        "path": input_path,
                        "message": "[%s/%s] 开始处理 %s"
                        % (file_index, file_count, os.path.basename(input_path)),
                    }
                )
                try:
                    result = batch.separate_file(input_path)
                    if use_dml:
                        message = "%s -> 成功 | %s | 推理 %.2fs | 编码 %.2fs" % (
                            os.path.basename(input_path),
                            batch.model_dtype.upper(),
                            result["inference_seconds"],
                            result["encode_seconds"],
                        )
                    else:
                        message = "%s -> 成功 | 推理 %.2fs | 编码 %.2fs" % (
                            os.path.basename(input_path),
                            result["inference_seconds"],
                            result["encode_seconds"],
                        )
                    successful_files += 1
                    _worker_emit(
                        {
                            "event": "file",
                            "ok": True,
                            "file_index": file_index,
                            "file_count": file_count,
                            "path": input_path,
                            "message": message,
                        }
                    )
                except Exception as error:
                    if (
                        use_dml
                        and successful_files == 0
                        and _fp16_retryable(
                            error, requested_dtype, allow_fp32_retry
                        )
                    ):
                        _worker_emit(
                            {
                                "event": "retry_fp32",
                                "message": "DirectML FP16 路径不兼容，准备改用 FP32",
                                "detail": traceback.format_exc(),
                                "file_index": file_index,
                                "file_count": file_count,
                            }
                        )
                        return 75
                    failed_files += 1
                    _worker_emit(
                        {
                            "event": "file",
                            "ok": False,
                            "file_index": file_index,
                            "file_count": file_count,
                            "path": input_path,
                            "message": "%s -> 失败\n%s"
                            % (os.path.basename(input_path), traceback.format_exc()),
                        }
                    )

        _worker_emit(
            {
                "event": "done",
                "processed": len(input_paths),
                "file_count": file_count,
                "successful": successful_files,
                "failed": failed_files,
                "model_dtype": requested_dtype,
            }
        )
        return 0
    except BaseException as error:
        try:
            retry_fp32 = use_dml and _fp16_retryable(
                error, requested_dtype, allow_fp32_retry
            )
            _worker_emit(
                {
                    "event": "fatal",
                    "message": "%s子进程失败\n%s"
                    % ("DML " if use_dml else "PyMSS ", traceback.format_exc()),
                    "retry_fp32": retry_fp32,
                }
            )
        except Exception:
            pass
        return 1


def _pymss_dml_worker_main(request_path):
    return _pymss_worker_main(request_path)


def _begin_pymss_task():
    with PYMSS_WORKER_STATE_LOCK:
        if PYMSS_WORKER_STATE["active"]:
            return None
        PYMSS_WORKER_STATE["task_id"] += 1
        PYMSS_WORKER_STATE["active"] = True
        PYMSS_WORKER_STATE["process"] = None
        PYMSS_WORKER_STATE["stop_requested"] = False
        return PYMSS_WORKER_STATE["task_id"]


def _finish_pymss_task(task_id):
    with PYMSS_WORKER_STATE_LOCK:
        if PYMSS_WORKER_STATE["task_id"] != task_id:
            return
        PYMSS_WORKER_STATE["active"] = False
        PYMSS_WORKER_STATE["process"] = None
        PYMSS_WORKER_STATE["stop_requested"] = False


def _pymss_task_stop_requested(task_id):
    with PYMSS_WORKER_STATE_LOCK:
        return (
            PYMSS_WORKER_STATE["task_id"] == task_id
            and PYMSS_WORKER_STATE["stop_requested"]
        )


def _register_worker(task_id, process):
    with PYMSS_WORKER_STATE_LOCK:
        if (
            PYMSS_WORKER_STATE["task_id"] != task_id
            or not PYMSS_WORKER_STATE["active"]
        ):
            return True
        PYMSS_WORKER_STATE["process"] = process
        return PYMSS_WORKER_STATE["stop_requested"]


def _unregister_worker(task_id, process):
    with PYMSS_WORKER_STATE_LOCK:
        if (
            PYMSS_WORKER_STATE["task_id"] == task_id
            and PYMSS_WORKER_STATE["process"] is process
        ):
            PYMSS_WORKER_STATE["process"] = None


def stop_pymss_separation():
    with PYMSS_WORKER_STATE_LOCK:
        if not PYMSS_WORKER_STATE["active"]:
            return "当前没有正在运行的 PyMSS 任务。"
        PYMSS_WORKER_STATE["stop_requested"] = True
        process = PYMSS_WORKER_STATE["process"]
    if process is not None:
        kill_process_tree(process, "PyMSS", logger)
    return "已请求停止 PyMSS 分离任务。"


def _read_worker_log(path, limit=12000):
    try:
        with open(path, "rb") as log_file:
            log_file.seek(0, os.SEEK_END)
            size = log_file.tell()
            log_file.seek(max(0, size - limit), os.SEEK_SET)
            return log_file.read().decode("utf-8", errors="replace").strip()
    except OSError:
        return ""


def _pymss_worker_events(
    task_id,
    spec,
    input_paths,
    desired_root,
    secondary_root,
    output_format,
    model_dtype,
    allow_fp32_retry,
):
    request = {
        "model_id": spec.model_id,
        "input_paths": input_paths,
        "desired_root": desired_root,
        "secondary_root": secondary_root,
        "output_format": output_format,
        "model_dtype": model_dtype,
        "allow_fp32_retry": allow_fp32_retry,
    }
    worker_code = (
        "import sys;"
        "request_path=sys.argv[1];"
        "sys.argv[:]=[sys.argv[0]];"
        "from tools.pymss_webui import _pymss_worker_main;"
        "raise SystemExit(_pymss_worker_main(request_path))"
    )

    if _pymss_task_stop_requested(task_id):
        yield {
            "event": "cancelled",
            "message": "PyMSS 分离任务已停止。",
            "file_count": len(input_paths),
        }
        return

    with tempfile.TemporaryDirectory(prefix="rvc-pymss-") as temp_root:
        request_path = Path(temp_root) / "request.json"
        stderr_path = Path(temp_root) / "worker.stderr.log"
        request_path.write_text(
            json.dumps(request, ensure_ascii=False),
            encoding="utf-8",
        )

        process = None
        return_code = None
        saw_done = False
        saw_fatal = False
        retry_event = None
        cancel_requested = False
        protocol_noise = deque(maxlen=5)
        with open(stderr_path, "wb") as stderr_file:
            try:
                environment = os.environ.copy()
                environment["PYTHONIOENCODING"] = "utf-8"
                environment["PYTHONUTF8"] = "1"
                process_group_kwargs = {}
                if os.name == "nt":
                    process_group_kwargs["creationflags"] = (
                        subprocess.CREATE_NO_WINDOW
                        | subprocess.CREATE_NEW_PROCESS_GROUP
                    )
                else:
                    process_group_kwargs["start_new_session"] = True
                process = subprocess.Popen(
                    [
                        sys.executable,
                        "-u",
                        "-c",
                        worker_code,
                        str(request_path),
                    ],
                    cwd=str(PROJECT_ROOT),
                    env=environment,
                    stdin=subprocess.DEVNULL,
                    stdout=subprocess.PIPE,
                    stderr=stderr_file,
                    text=True,
                    encoding="utf-8",
                    errors="replace",
                    bufsize=1,
                    **process_group_kwargs,
                )
                if _register_worker(task_id, process):
                    kill_process_tree(process, "PyMSS", logger)
                for raw_line in process.stdout:
                    line = raw_line.strip()
                    if not line:
                        continue
                    try:
                        event = json.loads(line)
                    except json.JSONDecodeError:
                        protocol_noise.append(line)
                        continue
                    if not isinstance(event, dict):
                        protocol_noise.append(line)
                        continue
                    event_type = event.get("event")
                    if event_type == "done":
                        saw_done = True
                        logger.info(
                            "PyMSS worker completed %s file(s)",
                            event.get("processed", 0),
                        )
                        yield event
                    elif event_type == "retry_fp32":
                        retry_event = event
                        logger.warning("PyMSS worker: %s", event.get("message", ""))
                    elif event_type == "fatal":
                        if event.get("retry_fp32"):
                            retry_event = event
                        else:
                            saw_fatal = True
                            logger.error("PyMSS worker: %s", event.get("message", ""))
                            yield event
                    elif event_type == "log":
                        log_level = getattr(
                            logging,
                            str(event.get("level", "INFO")).upper(),
                            logging.INFO,
                        )
                        logger.log(log_level, "PyMSS worker: %s", event.get("message", ""))
                        yield event
                    elif event_type == "file":
                        log_method = logger.info if event.get("ok") else logger.error
                        log_method("PyMSS worker: %s", event.get("message", ""))
                        yield event
                    elif event_type in {"file_start", "status"}:
                        logger.info("PyMSS worker: %s", event.get("message", ""))
                        yield event
                    elif event_type == "progress":
                        yield event
                    else:
                        protocol_noise.append(line)
                return_code = process.wait()
            finally:
                cancel_requested = _pymss_task_stop_requested(task_id)
                if process is not None:
                    try:
                        kill_process_tree(process, "PyMSS", logger)
                    finally:
                        try:
                            if process.stdout is not None:
                                process.stdout.close()
                        finally:
                            _unregister_worker(task_id, process)

        if protocol_noise:
            logger.warning(
                "Ignored non-protocol PyMSS worker stdout: %s",
                " | ".join(protocol_noise),
            )
        if cancel_requested:
            yield {
                "event": "cancelled",
                "message": "PyMSS 分离任务已停止。",
                "file_count": len(input_paths),
            }
            return
        if retry_event is not None:
            detail = retry_event.get("detail") or retry_event.get("message") or ""
            logger.warning("DirectML FP16 worker requested FP32 fallback: %s", detail)
            raise DMLFP16Fallback(detail)
        if saw_fatal:
            return
        if return_code != 0 or not saw_done:
            detail = _read_worker_log(stderr_path)
            if protocol_noise:
                noise = "\n".join(protocol_noise)
                detail = (
                    "%s\nstdout:\n%s" % (detail, noise)
                    if detail
                    else "stdout:\n%s" % noise
                )
            raise RuntimeError(
                "PyMSS worker exited unexpectedly (code=%s)%s"
                % (return_code, "\n%s" % detail if detail else "")
            )


def pymss_separate(
    model_name,
    inp_root,
    save_root_vocal,
    paths,
    save_root_ins,
    format0,
    event_callback=None,
):
    infos = deque(maxlen=240)
    task_id = None

    def notify(event):
        if event_callback is None:
            return
        try:
            event_callback(event)
        except Exception:
            logger.exception("PyMSS UI event callback failed")

    try:
        spec = resolve_model(model_name)
        input_paths = collect_input_paths(inp_root, paths)
        if not input_paths:
            raise ValueError("没有找到可处理的音频文件")

        task_id = _begin_pymss_task()
        if task_id is None:
            message = "已有 PyMSS 分离任务正在运行，请先停止当前任务。"
            infos.append(message)
            notify({"event": "busy", "message": message})
            yield "\n".join(infos)
            return

        device_type = torch.device(config.device).type
        use_dml = device_type == "privateuseone"
        infos.append("%s | %s | 正在加载模型" % (spec.label, spec.model_id))
        notify(
            {
                "event": "preparing",
                "message": "正在启动 PyMSS 模型子进程",
                "file_count": len(input_paths),
            }
        )
        logger.info(
            "Starting PyMSS worker: model=%s, device=%s, files=%s",
            spec.model_id,
            device_type,
            len(input_paths),
        )
        yield "\n".join(infos)

        attempts = _dml_model_dtype_attempts(spec) if use_dml else ("auto",)
        for attempt_index, model_dtype in enumerate(attempts):
            if _pymss_task_stop_requested(task_id):
                event = {
                    "event": "cancelled",
                    "message": "PyMSS 分离任务已停止。",
                    "file_count": len(input_paths),
                }
                notify(event)
                infos.append(event["message"])
                yield "\n".join(infos)
                return

            if use_dml:
                attempt_message = "DirectML 正在尝试参数精度 %s" % model_dtype.upper()
                infos.append(attempt_message)
                logger.info(attempt_message)
                notify(
                    {
                        "event": "precision_attempt",
                        "message": attempt_message,
                        "file_count": len(input_paths),
                    }
                )
                yield "\n".join(infos)

            events = _pymss_worker_events(
                task_id,
                spec,
                input_paths,
                save_root_vocal,
                save_root_ins,
                format0,
                model_dtype,
                allow_fp32_retry=(
                    use_dml
                    and model_dtype == "float16"
                    and attempt_index + 1 < len(attempts)
                ),
            )
            try:
                for event in events:
                    event_type = event.get("event")
                    notify(event)
                    if event_type == "progress":
                        if event_callback is not None:
                            yield "\n".join(infos)
                        continue
                    if event_type == "done":
                        if event_callback is not None:
                            yield "\n".join(infos)
                        continue

                    message = event.get("message")
                    if message:
                        infos.append(message)
                        yield "\n".join(infos)
                    if event_type == "cancelled":
                        return
            except DMLFP16Fallback:
                DML_FP16_DISABLED_MODEL_TYPES.add(spec.model_type)
                retry_message = (
                    "本机 DirectML 的 %s FP16 路径已降级；当前任务改用全新 FP32 子进程重试"
                    % spec.model_type
                )
                infos.append(retry_message)
                logger.warning(retry_message)
                notify(
                    {
                        "event": "retry_fp32",
                        "message": retry_message,
                        "file_count": len(input_paths),
                    }
                )
                yield "\n".join(infos)
                continue
            finally:
                events.close()
            break
    except Exception:
        if task_id is not None and _pymss_task_stop_requested(task_id):
            message = "PyMSS 分离任务已停止。"
            infos.append(message)
            notify({"event": "cancelled", "message": message})
        else:
            detail = traceback.format_exc()
            infos.append("失败\n%s" % detail)
            logger.error("PyMSS separation failed\n%s", detail)
            notify({"event": "fatal", "message": detail})
    finally:
        if task_id is not None:
            _finish_pymss_task(task_id)
    yield "\n".join(infos)