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, }