457 lines
17 KiB
Python
457 lines
17 KiB
Python
# demucs_pipeline.py
|
||
# Demucs 人声分离管道：FFmpeg 拉流 → Demucs 分离人声 → FFmpeg 推流
|
||
# 兼容 Windows / Ubuntu，支持 NVIDIA GPU 加速与 CPU 回退
|
||
"""
|
||
Demucs 音频管道模块
|
||
用于将混音分离为人声，推送到 YouTube 时避免 Content ID 识别背景音乐。
|
||
"""
|
||
import os
|
||
import sys
|
||
import socket
|
||
import threading
|
||
import subprocess
|
||
import queue
|
||
import time
|
||
import platform
|
||
|
||
try:
|
||
from stream_stats import StreamStats, parse_and_maybe_print
|
||
_STREAM_STATS_AVAILABLE = True
|
||
except ImportError:
|
||
_STREAM_STATS_AVAILABLE = False
|
||
|
||
# 可选依赖：Demucs 和 PyTorch
|
||
DEMUCS_AVAILABLE = False
|
||
DEMUCS_MODEL = None
|
||
TORCH_AVAILABLE = False
|
||
|
||
try:
|
||
import torch
|
||
TORCH_AVAILABLE = True
|
||
except ImportError:
|
||
pass
|
||
|
||
try:
|
||
from demucs.pretrained import get_model
|
||
from demucs.apply import apply_model
|
||
DEMUCS_AVAILABLE = True
|
||
except ImportError:
|
||
pass
|
||
|
||
|
||
def _check_demucs_available():
|
||
"""检查 Demucs 是否可用"""
|
||
return DEMUCS_AVAILABLE and TORCH_AVAILABLE
|
||
|
||
|
||
def _load_demucs_model(use_gpu: bool):
|
||
"""加载 Demucs 模型，支持 GPU/CPU 回退"""
|
||
global DEMUCS_MODEL
|
||
if not _check_demucs_available():
|
||
return None
|
||
try:
|
||
# htdemucs 较轻量，适合实时；htdemucs_ft 质量更好但更慢
|
||
model_name = "htdemucs"
|
||
device = "cpu"
|
||
if use_gpu and TORCH_AVAILABLE and torch.cuda.is_available():
|
||
try:
|
||
device = "cuda"
|
||
except Exception:
|
||
device = "cpu"
|
||
model = get_model(model_name)
|
||
model.eval()
|
||
return model.to(device)
|
||
except Exception as e:
|
||
print(f"[WARN] Demucs 模型加载失败，回退原始音频: {e}")
|
||
return None
|
||
|
||
|
||
# Demucs 模型采样率
|
||
DEMUCS_SAMPLE_RATE = 44100
|
||
DEMUCS_CHANNELS = 2
|
||
# 每次处理的音频长度（秒），1.0 更易实时，2.0 质量更好但 P40 可能跟不上
|
||
DEMUCS_CHUNK_SEC = 1.0
|
||
# 字节 per sample (s16le: 2 bytes)
|
||
BYTES_PER_SAMPLE = 2
|
||
|
||
|
||
def _demucs_separate_vocals(model, audio_chunk, device, vocal_mix_ratio=1.0, mix_source="bass"):
|
||
"""
|
||
对音频块进行人声分离，返回人声轨道。
|
||
vocal_mix_ratio: 1.0=纯人声；0.75=75%人声+25%混合（更自然）
|
||
mix_source: "bass"=与低音混合(少混响无重复感)；"other"=与其他混合(可能带混响)
|
||
"""
|
||
import torch
|
||
# audio_chunk: numpy or tensor, shape (channels, samples)
|
||
if hasattr(audio_chunk, 'numpy'):
|
||
wav = audio_chunk
|
||
else:
|
||
import numpy as np
|
||
wav = torch.from_numpy(audio_chunk).float()
|
||
# Demucs 期望 (batch, channels, samples)
|
||
if wav.dim() == 2:
|
||
wav = wav.unsqueeze(0)
|
||
with torch.no_grad():
|
||
out = apply_model(model, wav, device=device, progress=False, overlap=0.1)
|
||
# out: (batch, sources, channels, samples)
|
||
# sources: [drums=0, bass=1, other=2, vocals=3]
|
||
vocals = out[0, 3]
|
||
if vocal_mix_ratio >= 1.0:
|
||
return vocals.cpu().numpy()
|
||
# 混合 bass 或 other：bass 无混响、无重复感；other 可能带混响导致"重复一遍"
|
||
blend = out[0, 1] if mix_source == "bass" else out[0, 2]
|
||
mixed = vocals * vocal_mix_ratio + blend * (1.0 - vocal_mix_ratio)
|
||
return mixed.cpu().numpy()
|
||
|
||
|
||
def _run_demucs_worker(audio_read_pipe, audio_write_socket, model, device, stop_event, vocal_mix_ratio=1.0, mix_source="bass"):
|
||
"""Demucs 工作线程：从 FFmpeg 读音频，分离人声，写入 socket"""
|
||
import numpy as np
|
||
chunk_samples = int(DEMUCS_SAMPLE_RATE * DEMUCS_CHUNK_SEC * DEMUCS_CHANNELS)
|
||
chunk_bytes = chunk_samples * BYTES_PER_SAMPLE
|
||
buf = b""
|
||
|
||
def read_audio(size):
|
||
if hasattr(audio_read_pipe, 'recv'):
|
||
return audio_read_pipe.recv(size)
|
||
return audio_read_pipe.read(size)
|
||
|
||
try:
|
||
while not stop_event.is_set():
|
||
try:
|
||
data = read_audio(65536)
|
||
if not data:
|
||
break
|
||
buf += data
|
||
except (socket.error, OSError, ValueError):
|
||
break
|
||
while len(buf) >= chunk_bytes:
|
||
chunk = buf[:chunk_bytes]
|
||
buf = buf[chunk_bytes:]
|
||
# s16le -> float
|
||
arr = np.frombuffer(chunk, dtype=np.int16)
|
||
arr = arr.reshape(-1, DEMUCS_CHANNELS).T.astype(np.float32) / 32768.0
|
||
try:
|
||
vocals = _demucs_separate_vocals(model, arr, device, vocal_mix_ratio, mix_source)
|
||
# float -> s16le
|
||
out = (vocals * 32768).astype(np.int16).tobytes()
|
||
try:
|
||
audio_write_socket.sendall(out)
|
||
except (socket.error, OSError):
|
||
return
|
||
except Exception as e:
|
||
print(f"[WARN] Demucs 处理异常: {e}")
|
||
# 回退：直接透传原始音频
|
||
try:
|
||
audio_write_socket.sendall(chunk)
|
||
except (socket.error, OSError):
|
||
return
|
||
# 剩余 buffer
|
||
if buf and audio_write_socket:
|
||
try:
|
||
arr = np.frombuffer(buf, dtype=np.int16)
|
||
arr = arr.reshape(-1, DEMUCS_CHANNELS).T.astype(np.float32) / 32768.0
|
||
vocals = _demucs_separate_vocals(model, arr, device, vocal_mix_ratio, mix_source)
|
||
out = (vocals * 32768).astype(np.int16).tobytes()
|
||
audio_write_socket.sendall(out)
|
||
except Exception:
|
||
pass
|
||
except Exception as e:
|
||
print(f"[WARN] Demucs worker 异常: {e}")
|
||
finally:
|
||
try:
|
||
audio_write_socket.shutdown(socket.SHUT_WR)
|
||
audio_write_socket.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def run_demucs_pipeline(
|
||
real_url: str,
|
||
youtube_rtmp: str,
|
||
douyin_headers: str,
|
||
nvenc_params: dict,
|
||
end_keywords: list,
|
||
stderr_queue_timeout: float,
|
||
no_frame_timeout: int,
|
||
start_check_after: int,
|
||
stale_output_seconds: int,
|
||
stream_ended_exception,
|
||
cleanup_proc_fn,
|
||
stderr_reader_fn,
|
||
vocal_mix_ratio: float = 0.75,
|
||
mix_source: str = "bass",
|
||
stats_interval_sec: int = 10,
|
||
):
|
||
"""
|
||
运行 Demucs 管道：拉流 → 人声分离 → 推流
|
||
若主播下播应停止则 raise stream_ended_exception
|
||
"""
|
||
|
||
# 检查 Demucs
|
||
use_gpu = nvenc_params.get("demucs_use_gpu", True)
|
||
model = _load_demucs_model(use_gpu)
|
||
if model is None:
|
||
raise RuntimeError("Demucs 模型不可用，将回退原始推流")
|
||
|
||
device = next(model.parameters()).device
|
||
device_str = "cuda" if "cuda" in str(device) else "cpu"
|
||
src_name = "低音" if mix_source == "bass" else "其他"
|
||
mix_info = "纯人声" if vocal_mix_ratio >= 1.0 else f"{int(vocal_mix_ratio*100)}%人声+{int((1-vocal_mix_ratio)*100)}%{src_name}"
|
||
print(f"[INFO] Demucs 人声分离已启用 ({device_str})，模式: {mix_info}，低通 9kHz 降尖锐")
|
||
|
||
# TCP 服务器用于音频（跨平台）
|
||
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||
server.bind(("127.0.0.1", 0))
|
||
port = server.getsockname()[1]
|
||
server.listen(1)
|
||
|
||
tcp_audio_url = f"tcp://127.0.0.1:{port}"
|
||
stop_event = threading.Event()
|
||
conn = [None] # 用 list 以便在闭包中修改
|
||
|
||
def accept_conn():
|
||
server.settimeout(30)
|
||
try:
|
||
c, _ = server.accept()
|
||
conn[0] = c
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
try:
|
||
server.close()
|
||
except Exception:
|
||
pass
|
||
|
||
accept_thread = threading.Thread(target=accept_conn, daemon=True)
|
||
accept_thread.start()
|
||
time.sleep(0.2) # 确保 server 已 listen
|
||
|
||
# FFmpeg 拉流（视频）
|
||
ffmpeg_v_cmd = [
|
||
"ffmpeg", "-y", "-loglevel", "info", "-nostdin", "-re",
|
||
"-stats", "-stats_period", "1",
|
||
"-rw_timeout", "30000000",
|
||
"-reconnect", "1", "-reconnect_at_eof", "1", "-reconnect_streamed", "1",
|
||
"-reconnect_delay_max", "5",
|
||
"-fflags", "+genpts+discardcorrupt+nobuffer+flush_packets",
|
||
"-err_detect", "ignore_err",
|
||
"-headers", douyin_headers,
|
||
"-i", real_url,
|
||
"-map", "0:v", "-c:v", "copy", "-f", "nut", "pipe:1",
|
||
]
|
||
|
||
# FFmpeg 拉流（音频）
|
||
ffmpeg_a_cmd = [
|
||
"ffmpeg", "-y", "-loglevel", "error", "-nostdin", "-re",
|
||
"-rw_timeout", "30000000",
|
||
"-reconnect", "1", "-reconnect_at_eof", "1", "-reconnect_streamed", "1",
|
||
"-headers", douyin_headers,
|
||
"-i", real_url,
|
||
"-map", "0:a", "-f", "s16le", "-ac", str(DEMUCS_CHANNELS),
|
||
"-ar", str(DEMUCS_SAMPLE_RATE), "pipe:1",
|
||
]
|
||
|
||
# FFmpeg 推流：音频 tcp 放第一（FFmpeg 按顺序打开输入，pipe:0 会阻塞直到有数据，
|
||
# 若 pipe 在前则永远连不上 tcp，导致 accept 超时）
|
||
vf = "fps=30,scale=720:1280:force_original_aspect_ratio=decrease:flags=lanczos,pad=720:1280:(ow-iw)/2:(oh-ih)/2:black"
|
||
ffmpeg_out_cmd = [
|
||
"ffmpeg", "-y", "-loglevel", "info", "-nostdin",
|
||
"-f", "s16le", "-ac", str(DEMUCS_CHANNELS), "-ar", str(DEMUCS_SAMPLE_RATE),
|
||
"-i", tcp_audio_url,
|
||
"-f", "nut", "-i", "pipe:0",
|
||
"-map", "1:v", "-map", "0:a",
|
||
"-vf", vf,
|
||
"-c:v", nvenc_params.get("codec_v", "libx264"),
|
||
"-preset", nvenc_params.get("preset_v", "fast"),
|
||
"-profile:v", "high",
|
||
"-b:v", f"{nvenc_params.get('video_bitrate_k', 4500)}k",
|
||
"-maxrate", f"{nvenc_params.get('video_bitrate_k', 4500)+500}k",
|
||
"-bufsize", f"{nvenc_params.get('video_bitrate_k', 4500)*2}k",
|
||
"-fps_mode", "cfr", "-g", "60", "-keyint_min", "60", "-r", "30",
|
||
"-bf", "0",
|
||
"-pix_fmt", "yuv420p",
|
||
"-c:a", "aac", "-b:a", "128k", "-ar", "48000", "-ac", "2",
|
||
"-af", "lowpass=f=9000,aresample=async=1:first_pts=0",
|
||
"-flags", "+global_header", "-flvflags", "no_duration_filesize",
|
||
"-f", "flv", youtube_rtmp,
|
||
]
|
||
# NVENC 用 tune_param；libx264 用 sc_threshold + x264-params
|
||
codec = nvenc_params.get("codec_v", "libx264")
|
||
if nvenc_params.get("tune_param") and codec == "h264_nvenc":
|
||
idx = ffmpeg_out_cmd.index("-profile:v")
|
||
for p in reversed(nvenc_params["tune_param"]):
|
||
ffmpeg_out_cmd.insert(idx, p)
|
||
elif codec == "libx264":
|
||
idx = ffmpeg_out_cmd.index("-bf")
|
||
ffmpeg_out_cmd.insert(idx + 2, "-sc_threshold")
|
||
ffmpeg_out_cmd.insert(idx + 3, "0")
|
||
idx = ffmpeg_out_cmd.index("-pix_fmt")
|
||
ffmpeg_out_cmd.insert(idx, "-x264-params")
|
||
ffmpeg_out_cmd.insert(idx + 1, "force-cfr=1:scenecut=0:open_gop=0:sync-lookahead=0")
|
||
|
||
proc_v = None
|
||
proc_a = None
|
||
proc_out = None
|
||
demucs_thread = None
|
||
|
||
try:
|
||
proc_v = subprocess.Popen(
|
||
ffmpeg_v_cmd,
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
bufsize=1024 * 1024,
|
||
)
|
||
proc_a = subprocess.Popen(
|
||
ffmpeg_a_cmd,
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.DEVNULL,
|
||
bufsize=1024 * 1024,
|
||
)
|
||
proc_out = subprocess.Popen(
|
||
ffmpeg_out_cmd,
|
||
stdin=subprocess.PIPE,
|
||
stdout=subprocess.DEVNULL,
|
||
stderr=subprocess.PIPE,
|
||
bufsize=1024 * 1024,
|
||
)
|
||
|
||
# 等待 TCP 连接（FFmpeg_out 会连接）
|
||
accept_thread.join(timeout=15)
|
||
if conn[0] is None:
|
||
raise RuntimeError("Demucs 管道: 音频 TCP 连接超时")
|
||
|
||
# 视频转发线程
|
||
def copy_video():
|
||
try:
|
||
while proc_v.poll() is None and proc_out.poll() is None:
|
||
chunk = proc_v.stdout.read(1024 * 256)
|
||
if not chunk:
|
||
break
|
||
proc_out.stdin.write(chunk)
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
try:
|
||
proc_out.stdin.close()
|
||
except Exception:
|
||
pass
|
||
|
||
copy_t = threading.Thread(target=copy_video, daemon=True)
|
||
copy_t.start()
|
||
|
||
# Demucs 工作线程（从 proc_a 读，处理后写 conn）
|
||
# proc_a 输出是 stdout，conn 是 socket。需要把 proc_a.stdout 的数据通过 socket 发到 FFmpeg。
|
||
# FFmpeg_out 的第二个输入是 tcp，已经连接了 conn。所以我们从 proc_a 读，Demucs 处理，写 conn。
|
||
# 但 proc_a.stdout 是 pipe，conn 是 socket。我们需要一个线程读 proc_a 并写 conn。
|
||
# 注意：proc_a 输出到 stdout，我们读。conn 是 FFmpeg 作为 client 连上来的，所以 FFmpeg 在读，我们写 conn。
|
||
demucs_thread = threading.Thread(
|
||
target=_run_demucs_worker,
|
||
args=(proc_a.stdout, conn[0], model, device, stop_event, vocal_mix_ratio, mix_source),
|
||
daemon=True,
|
||
)
|
||
demucs_thread.start()
|
||
|
||
# 监控 proc_v 的 stderr（与源连接，用于检测下播）
|
||
stderr_q = queue.Queue()
|
||
reader_t = threading.Thread(target=stderr_reader_fn, args=(proc_v, stderr_q))
|
||
reader_t.daemon = True
|
||
reader_t.start()
|
||
|
||
# 监控 proc_out 的 stderr（推流到 YouTube），解析状态用于详细日志
|
||
stream_stats = None
|
||
if _STREAM_STATS_AVAILABLE:
|
||
stream_stats = StreamStats(interval_sec=stats_interval_sec)
|
||
stream_stats.start_periodic_printer()
|
||
print(f"[INFO] 推流状态监控已启用，每 {stats_interval_sec}s 输出帧率/码率/网络等")
|
||
|
||
def _proc_out_stderr_reader():
|
||
try:
|
||
for line in iter(proc_out.stderr.readline, b""):
|
||
if line:
|
||
decoded = line.decode("utf-8", errors="replace").strip()
|
||
if decoded:
|
||
if stream_stats and parse_and_maybe_print(decoded, stream_stats, "[YouTube推流]", verbose=False):
|
||
pass # 已由 parse 处理，进度行不重复打印
|
||
else:
|
||
print(f"[YouTube推流] {decoded}")
|
||
except Exception:
|
||
pass
|
||
|
||
proc_out_reader_t = threading.Thread(target=_proc_out_stderr_reader, daemon=True)
|
||
proc_out_reader_t.start()
|
||
|
||
start_time = last_frame_time = time.time()
|
||
proc_v.last_out_ts = time.time()
|
||
|
||
while True:
|
||
try:
|
||
line = stderr_q.get(timeout=stderr_queue_timeout)
|
||
except queue.Empty:
|
||
# 若 proc_out（推流进程）已异常退出，立即报告
|
||
if proc_out.poll() is not None and proc_out.returncode != 0:
|
||
print(f"[ERROR] YouTube 推流进程已退出，返回码: {proc_out.returncode}，可能是 Demucs 处理跟不上或网络问题")
|
||
print(f"[TIP] 若频繁断流，可尝试 config/youtube.ini 中设置 enable_demucs = 否 使用原始音频")
|
||
break
|
||
if time.time() - getattr(proc_v, "last_out_ts", time.time()) > stale_output_seconds:
|
||
print("[WARN] 长时间无任何日志，强制重启进程")
|
||
break
|
||
continue
|
||
|
||
if line is None:
|
||
if proc_v.poll() is not None:
|
||
break
|
||
continue
|
||
|
||
# proc_v 的 stderr 为 bytes，需解码为 str
|
||
if isinstance(line, bytes):
|
||
line = line.decode("utf-8", errors="replace")
|
||
line = line.strip()
|
||
if line.startswith("__ERR__READER__"):
|
||
print(f"[ERROR] stderr reader 异常: {line}")
|
||
continue
|
||
|
||
if line:
|
||
proc_v.last_out_ts = time.time()
|
||
if stream_stats:
|
||
stream_stats.parse_lag_event(line) # 源端抖动统计
|
||
print(f"[FFmpeg] {line}")
|
||
line_lower = line.lower()
|
||
if any(kw in line_lower for kw in end_keywords):
|
||
print("[ERROR] 检测到主播真下播或被明确拒绝，停止推流")
|
||
raise stream_ended_exception()
|
||
if "frame=" in line_lower or "fps=" in line_lower:
|
||
last_frame_time = time.time()
|
||
if time.time() - start_time > start_check_after:
|
||
if time.time() - last_frame_time > no_frame_timeout:
|
||
print(f"[WARN] 连续 {no_frame_timeout}s 无新帧，判定主播已下播")
|
||
raise stream_ended_exception()
|
||
|
||
except stream_ended_exception:
|
||
raise
|
||
except Exception as e:
|
||
print(f"[ERROR] Demucs 管道异常: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
finally:
|
||
stop_event.set()
|
||
if stream_stats is not None:
|
||
stream_stats.stop()
|
||
if proc_out is not None and proc_out.poll() is not None and proc_out.returncode != 0:
|
||
print(f"[INFO] YouTube 推流进程退出码: {proc_out.returncode}")
|
||
cleanup_proc_fn(proc_v)
|
||
cleanup_proc_fn(proc_a)
|
||
cleanup_proc_fn(proc_out)
|
||
try:
|
||
if conn[0]:
|
||
conn[0].close()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
server.close()
|
||
except Exception:
|
||
pass
|