194 lines
7.7 KiB
Python
194 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
import uuid
|
|
import wave
|
|
from collections import deque
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from config import settings
|
|
|
|
|
|
class AvatarService:
|
|
def __init__(self, width: int = 640, height: int = 360) -> None:
|
|
self.width = width
|
|
self.height = height
|
|
self._frames: deque[np.ndarray] = deque()
|
|
self._lock = asyncio.Lock()
|
|
self._job_lock = asyncio.Lock()
|
|
self._last_error: str | None = None
|
|
self._base_frame = self._load_base_frame()
|
|
|
|
async def render_frame(self, mouth_open: float = 0.0, speaking: bool = False) -> np.ndarray:
|
|
frame = self._base_frame.copy()
|
|
if speaking:
|
|
alpha = float(np.clip(mouth_open, 0.0, 1.0))
|
|
# Slightly boost contrast/brightness while speaking to keep a "live" feeling.
|
|
frame = cv2.convertScaleAbs(frame, alpha=1.0 + 0.10 * alpha, beta=2 + int(8 * alpha))
|
|
# Fallback mouth animation on realistic base frame.
|
|
h, w = frame.shape[:2]
|
|
cx = w // 2
|
|
cy = int(h * 0.70)
|
|
half_w = max(28, int(w * 0.055))
|
|
half_h = max(4, int(4 + alpha * h * 0.035))
|
|
|
|
lip_color = (26, 26, 92)
|
|
lip_border = (70, 70, 160)
|
|
cv2.ellipse(frame, (cx, cy), (half_w + 2, half_h + 2), 0, 0, 360, lip_border, -1)
|
|
cv2.ellipse(frame, (cx, cy), (half_w, half_h), 0, 0, 360, lip_color, -1)
|
|
return frame
|
|
|
|
async def idle_frame(self) -> np.ndarray:
|
|
return await self.render_frame(mouth_open=0.0, speaking=False)
|
|
|
|
async def speaking_frame(self, mouth_open: float = 0.5) -> np.ndarray:
|
|
return await self.render_frame(mouth_open=mouth_open, speaking=True)
|
|
|
|
async def pop_generated_frame(self) -> np.ndarray | None:
|
|
async with self._lock:
|
|
if self._frames:
|
|
return self._frames.popleft()
|
|
return None
|
|
|
|
async def enqueue_musetalk_from_audio(self, audio: np.ndarray, sr: int) -> None:
|
|
if not settings.musetalk_enabled:
|
|
return
|
|
if audio.size <= 1 or sr <= 0:
|
|
return
|
|
|
|
async with self._job_lock:
|
|
try:
|
|
job_id = f"job_{uuid.uuid4().hex[:10]}"
|
|
repo_dir = Path(settings.musetalk_repo_dir)
|
|
work_root = repo_dir / "results" / "visual-chat-jobs" / job_id
|
|
work_root.mkdir(parents=True, exist_ok=True)
|
|
wav_path = work_root / "input.wav"
|
|
cfg_path = work_root / "inference.yaml"
|
|
output_name = f"{job_id}.mp4"
|
|
output_path = Path(settings.musetalk_result_dir) / "v15" / output_name
|
|
|
|
self._write_wav(wav_path, audio.astype(np.float32, copy=False), sr)
|
|
cfg_path.write_text(
|
|
(
|
|
"task_0:\n"
|
|
f" video_path: \"{settings.musetalk_source_video}\"\n"
|
|
f" audio_path: \"{wav_path}\"\n"
|
|
f" result_name: \"{output_name}\"\n"
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
env = os.environ.copy()
|
|
env.setdefault("HF_HOME", str(repo_dir / "models"))
|
|
env["PYTHONPATH"] = (
|
|
f"{repo_dir}:{env['PYTHONPATH']}" if env.get("PYTHONPATH") else str(repo_dir)
|
|
)
|
|
proc = await asyncio.create_subprocess_exec(
|
|
sys.executable,
|
|
str(settings.musetalk_infer_script),
|
|
"--version",
|
|
"v15",
|
|
"--inference_config",
|
|
str(cfg_path),
|
|
"--result_dir",
|
|
str(settings.musetalk_result_dir),
|
|
"--unet_model_path",
|
|
str(settings.musetalk_unet_model),
|
|
"--unet_config",
|
|
str(settings.musetalk_unet_config),
|
|
"--whisper_dir",
|
|
str(settings.musetalk_whisper_dir),
|
|
"--batch_size",
|
|
"8",
|
|
"--use_float16",
|
|
cwd=str(repo_dir),
|
|
env=env,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
_, stderr = await proc.communicate()
|
|
if proc.returncode != 0:
|
|
self._last_error = stderr.decode("utf-8", errors="ignore")[-500:]
|
|
return
|
|
frames: list[np.ndarray] = []
|
|
if output_path.exists():
|
|
frames = self._read_video_frames(str(output_path))
|
|
if not frames:
|
|
# Some MuseTalk runs leave image frames instead of final mp4.
|
|
fallback_dir = Path(settings.musetalk_result_dir) / "v15" / "yongen"
|
|
if fallback_dir.exists():
|
|
frames = self._read_image_frames(str(fallback_dir))
|
|
if not frames:
|
|
self._last_error = f"MuseTalk produced no readable frames: {output_path}"
|
|
return
|
|
|
|
async with self._lock:
|
|
self._frames.clear()
|
|
self._frames.extend(frames)
|
|
self._last_error = None
|
|
except Exception as exc: # pragma: no cover
|
|
self._last_error = str(exc)
|
|
|
|
def _read_video_frames(self, path: str) -> list[np.ndarray]:
|
|
cap = cv2.VideoCapture(path)
|
|
frames: list[np.ndarray] = []
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok or frame is None:
|
|
break
|
|
if frame.shape[1] != self.width or frame.shape[0] != self.height:
|
|
frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_LINEAR)
|
|
frames.append(frame)
|
|
cap.release()
|
|
return frames
|
|
|
|
def _read_image_frames(self, dir_path: str) -> list[np.ndarray]:
|
|
frames: list[np.ndarray] = []
|
|
for name in sorted(os.listdir(dir_path)):
|
|
if not name.lower().endswith(".png"):
|
|
continue
|
|
frame = cv2.imread(os.path.join(dir_path, name))
|
|
if frame is None:
|
|
continue
|
|
if frame.shape[1] != self.width or frame.shape[0] != self.height:
|
|
frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_LINEAR)
|
|
frames.append(frame)
|
|
return frames
|
|
|
|
def _load_base_frame(self) -> np.ndarray:
|
|
src = settings.musetalk_source_video
|
|
if src and os.path.exists(src):
|
|
cap = cv2.VideoCapture(src)
|
|
ok, frame = cap.read()
|
|
cap.release()
|
|
if ok and frame is not None:
|
|
if frame.shape[1] != self.width or frame.shape[0] != self.height:
|
|
frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_LINEAR)
|
|
return frame
|
|
# Last-resort dark background (should rarely happen).
|
|
fallback = np.zeros((self.height, self.width, 3), dtype=np.uint8)
|
|
fallback[:, :, :] = (20, 20, 20)
|
|
return fallback
|
|
|
|
@staticmethod
|
|
def _write_wav(path: Path, audio: np.ndarray, sr: int) -> None:
|
|
pcm = np.clip(audio * 32767.0, -32768, 32767).astype(np.int16)
|
|
with wave.open(str(path), "wb") as wf:
|
|
wf.setnchannels(1)
|
|
wf.setsampwidth(2)
|
|
wf.setframerate(sr)
|
|
wf.writeframes(pcm.tobytes())
|
|
|
|
@property
|
|
def health(self) -> dict:
|
|
return {
|
|
"enabled": settings.musetalk_enabled,
|
|
"queued_frames": len(self._frames),
|
|
"last_error": self._last_error,
|
|
}
|