PLUS: мультипроход со слиянием и проверкой нот по аудио

plus/serve_muscriptor.py

@@ -90,7 +90,31 @@ def healthy(port: int, timeout: float = 3.0) -> bool:
         return False
 
 
-def start_server(port: int, model: str, device: str, dtype: str | None) -> subprocess.Popen:
+def low_memory_env(threads: int | None = None) -> dict:
+    """Убрать раздувание памяти аллокаторами, не трогая точность модели.
+
+    Веса остаются float32, никакого квантования: экономия идёт за счёт числа
+    потоков (у каждого своя арена malloc) и возврата освобождённых страниц ОС.
+    """
+    cores = os.cpu_count() or 4
+    count = threads or max(1, min(4, cores // 2))
+    env = dict(os.environ)
+    env.update({
+        "OMP_NUM_THREADS": str(count),
+        "MKL_NUM_THREADS": str(count),
+        "OPENBLAS_NUM_THREADS": str(count),
+        "NUMEXPR_NUM_THREADS": str(count),
+        "MALLOC_ARENA_MAX": "2",
+        "MALLOC_TRIM_THRESHOLD_": "1000000",
+        "MALLOC_MMAP_THRESHOLD_": "131072",
+        "TOKENIZERS_PARALLELISM": "false",
+        "PYTORCH_CPU_ALLOC_CONF": "garbage_collection_threshold:0.6",
+    })
+    return env
+
+
+def start_server(port: int, model: str, device: str, dtype: str | None,
+                 env: dict | None = None) -> subprocess.Popen:
     command = [sys.executable, "-m", "muscriptor", "serve",
                "--model", model, "--device", device, "--host", "0.0.0.0", "--port", str(port)]
     if dtype:
@@ -98,7 +122,8 @@ def start_server(port: int, model: str, device: str, dtype: str | None) -> subpr
     log = open("muscriptor_server.log", "a")
     say("запускаю: " + " ".join(command[1:]))
     process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT,
-                               stdin=open(os.devnull, "rb"))
+                               stdin=open(os.devnull, "rb"),
+                               env=env or dict(os.environ))
     deadline = time.time() + 600
     started = time.time()
     while time.time() < deadline:
@@ -342,7 +367,8 @@ def ensure_audio_tools() -> bool:
 
 
 def start_proxy(server_port: int, listen_port: int,
-                instruments: str = BAND_INSTRUMENTS) -> subprocess.Popen:
+                instruments: str = BAND_INSTRUMENTS,
+                multi_pass: str = "") -> subprocess.Popen:
     """SSE-прокси: держит поток /transcribe живым через медленный туннель."""
     script = pathlib.Path(__file__).with_name("sse_proxy.py")
     if not script.exists():
@@ -350,7 +376,8 @@ def start_proxy(server_port: int, listen_port: int,
     proc = subprocess.Popen(
         [sys.executable, "-u", str(script),
          "--listen", str(listen_port), "--target", str(server_port),
-         "--instruments", instruments],
+         "--instruments", instruments]
+        + (["--multi-pass", multi_pass] if multi_pass else []),
         stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
     for _ in range(20):
         time.sleep(0.5)
@@ -391,8 +418,16 @@ def main(argv: list[str] | None = None) -> int:
                         help="имя пользователя для локального туннеля (по умолчанию nokey)")
     parser.add_argument("--no-tunnel", action="store_true", help="без публичной ссылки")
     parser.add_argument("--keep", type=int, default=0, help="держать N секунд и выйти")
+    parser.add_argument("--low-memory", action="store_true",
+                        help="экономный режим для слабого ПК: меньше потоков и арен,"
+                             " веса всё так же float32")
+    parser.add_argument("--threads", type=int, default=None,
+                        help="сколько потоков дать модели (по умолчанию половина ядер)")
     parser.add_argument("--skip-audio", action="store_true",
                         help="не ставить fluidsynth (рендер аудио будет недоступен)")
+    parser.add_argument("--multi-pass", default="",
+                        help="второй проход с другим набором инструментов и слияние,"
+                             " например: chromatic_percussion,electric_bass,acoustic_bass")
     parser.add_argument("--instruments", default=BAND_INSTRUMENTS,
                         help="инструменты для модели через запятую; auto = пусть решает сама")
     args = parser.parse_args(argv)
@@ -407,10 +442,15 @@ def main(argv: list[str] | None = None) -> int:
         ensure_audio_tools()
     ensure_ssh()
     ensure_key()
-    server = start_server(args.port, args.model, args.device, args.dtype)
+    env = None
+    if args.low_memory or args.threads:
+        env = low_memory_env(args.threads)
+        say("      экономный режим: потоков %s, арен malloc 2, веса float32"
+            % env["OMP_NUM_THREADS"])
+    server = start_server(args.port, args.model, args.device, args.dtype, env)
     say("[1/2] сервер отвечает: http://127.0.0.1:%d" % args.port)
     proxy_port = args.port + 1
-    proxy = start_proxy(args.port, proxy_port, args.instruments)
+    proxy = start_proxy(args.port, proxy_port, args.instruments, args.multi_pass)
     if proxy.poll() is None:
         say("      поток защищён прокси: http://127.0.0.1:%d" % proxy_port)
     else:
@@ -445,10 +485,11 @@ def main(argv: list[str] | None = None) -> int:
                 break
             if server.poll() is not None:
                 say("сервер упал, поднимаю заново")
-                server = start_server(args.port, args.model, args.device, args.dtype)
+                server = start_server(args.port, args.model, args.device, args.dtype, env)
             if proxy.poll() is not None:
                 say("прокси упал, поднимаю заново")
-                proxy = start_proxy(args.port, proxy_port, args.instruments)
+                proxy = start_proxy(args.port, proxy_port, args.instruments,
+                                    args.multi_pass)
             if tunnel is not None and not tunnel.alive():
                 say("туннель умер, открываю новый")
                 tunnel.kill()