init
This commit is contained in:
@@ -0,0 +1,193 @@
|
||||
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,
|
||||
}
|
||||
Reference in New Issue
Block a user