大功告成
This commit is contained in:
+190
-170
@@ -1,193 +1,213 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
import wave
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
import math
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
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()
|
||||
def __init__(self) -> None:
|
||||
self.driver_name = "arkit-blendshape-v1"
|
||||
self._lock = asyncio.Lock()
|
||||
self._job_lock = asyncio.Lock()
|
||||
self._sequence = 0
|
||||
self._last_error: str | None = None
|
||||
self._base_frame = self._load_base_frame()
|
||||
self._last_frame_count = 0
|
||||
self._last_chunk_duration_ms = 0
|
||||
|
||||
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))
|
||||
async def build_idle_payload(self, reason: str = "idle") -> dict[str, Any]:
|
||||
frame = self._build_frame(seq=await self._next_sequence(), time_ms=0, controls=self._neutral_controls())
|
||||
self._last_frame_count = 1
|
||||
self._last_chunk_duration_ms = 0
|
||||
self._last_error = None
|
||||
return {
|
||||
"type": "animation_state",
|
||||
"driver": self.driver_name,
|
||||
"mode": settings.avatar_driver_mode,
|
||||
"schema": settings.avatar_blendshape_schema,
|
||||
"fps": settings.avatar_fps,
|
||||
"reason": reason,
|
||||
"frame_count": 1,
|
||||
"frames": [frame],
|
||||
"ts_ms": int(time.time() * 1000),
|
||||
}
|
||||
|
||||
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 build_reset_payload(self, reason: str = "reset") -> dict[str, Any]:
|
||||
payload = await self.build_idle_payload(reason=reason)
|
||||
payload["type"] = "animation_reset"
|
||||
return payload
|
||||
|
||||
async def idle_frame(self) -> np.ndarray:
|
||||
return await self.render_frame(mouth_open=0.0, speaking=False)
|
||||
async def build_controls_from_audio(
|
||||
self,
|
||||
audio: np.ndarray,
|
||||
sr: int,
|
||||
*,
|
||||
text: str = "",
|
||||
chunk_index: int = 0,
|
||||
total_chunks: int = 1,
|
||||
) -> dict[str, Any]:
|
||||
if sr <= 0:
|
||||
self._last_error = f"invalid sample rate: {sr}"
|
||||
return await self.build_idle_payload(reason="invalid-sample-rate")
|
||||
|
||||
async def speaking_frame(self, mouth_open: float = 0.5) -> np.ndarray:
|
||||
return await self.render_frame(mouth_open=mouth_open, speaking=True)
|
||||
mono = np.asarray(audio, dtype=np.float32).reshape(-1)
|
||||
if mono.size <= 1:
|
||||
return await self.build_idle_payload(reason="empty-audio")
|
||||
|
||||
async def pop_generated_frame(self) -> np.ndarray | None:
|
||||
async with self._lock:
|
||||
if self._frames:
|
||||
return self._frames.popleft()
|
||||
return None
|
||||
mono = np.clip(mono, -1.0, 1.0)
|
||||
frame_count = max(1, int(math.ceil((mono.shape[0] / float(sr)) * settings.avatar_fps)))
|
||||
seqs = await self._reserve_sequences(frame_count)
|
||||
frames = self._audio_to_frames(mono, sr, seqs)
|
||||
|
||||
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())
|
||||
self._last_frame_count = len(frames)
|
||||
self._last_chunk_duration_ms = int(round(1000.0 * mono.shape[0] / float(sr)))
|
||||
self._last_error = None
|
||||
return {
|
||||
"type": "animation_chunk",
|
||||
"driver": self.driver_name,
|
||||
"mode": settings.avatar_driver_mode,
|
||||
"schema": settings.avatar_blendshape_schema,
|
||||
"chunk_index": chunk_index,
|
||||
"total_chunks": total_chunks,
|
||||
"sample_rate": sr,
|
||||
"fps": settings.avatar_fps,
|
||||
"text": text,
|
||||
"frame_count": len(frames),
|
||||
"duration_ms": self._last_chunk_duration_ms,
|
||||
"frames": frames,
|
||||
"ts_ms": int(time.time() * 1000),
|
||||
}
|
||||
|
||||
@property
|
||||
def health(self) -> dict:
|
||||
def schema(self) -> dict[str, Any]:
|
||||
return {
|
||||
"enabled": settings.musetalk_enabled,
|
||||
"queued_frames": len(self._frames),
|
||||
"driver": self.driver_name,
|
||||
"mode": settings.avatar_driver_mode,
|
||||
"protocol": settings.avatar_control_protocol,
|
||||
"schema": settings.avatar_blendshape_schema,
|
||||
"controls": [
|
||||
"jawOpen",
|
||||
"mouthClose",
|
||||
"mouthFunnel",
|
||||
"mouthPucker",
|
||||
"viseme_aa",
|
||||
"viseme_ee",
|
||||
"viseme_oh",
|
||||
"headYaw",
|
||||
"headPitch",
|
||||
"headRoll",
|
||||
],
|
||||
}
|
||||
|
||||
def _audio_to_frames(self, audio: np.ndarray, sr: int, sequences: range) -> list[dict[str, Any]]:
|
||||
samples_per_frame = max(1, int(round(sr / max(1, settings.avatar_fps))))
|
||||
frames: list[dict[str, Any]] = []
|
||||
jaw_prev = 0.0
|
||||
for index, seq in enumerate(sequences):
|
||||
start = index * samples_per_frame
|
||||
end = min(audio.shape[0], start + samples_per_frame)
|
||||
window = audio[start:end]
|
||||
if window.size == 0:
|
||||
window = np.zeros(1, dtype=np.float32)
|
||||
energy, brightness = self._extract_features(window, sr)
|
||||
jaw_open = 0.70 * jaw_prev + 0.30 * energy
|
||||
jaw_prev = jaw_open
|
||||
phase = index / max(1, len(sequences) - 1)
|
||||
frames.append(
|
||||
self._build_frame(
|
||||
seq=seq,
|
||||
time_ms=int(round(index * 1000.0 / settings.avatar_fps)),
|
||||
controls=self._controls_from_features(jaw_open, brightness, phase),
|
||||
)
|
||||
)
|
||||
return frames
|
||||
|
||||
def _extract_features(self, window: np.ndarray, sr: int) -> tuple[float, float]:
|
||||
rms = float(np.sqrt(np.mean(window**2))) if window.size else 0.0
|
||||
energy = float(np.clip((rms - 0.01) * 10.0, 0.0, 1.0))
|
||||
if window.size < 16:
|
||||
return energy, 0.5
|
||||
spectrum = np.abs(np.fft.rfft(window * np.hanning(window.size)))
|
||||
if spectrum.size <= 1:
|
||||
return energy, 0.5
|
||||
freqs = np.fft.rfftfreq(window.size, d=1.0 / float(sr))
|
||||
weights = spectrum[1:]
|
||||
freq_values = freqs[1:]
|
||||
denom = float(np.sum(weights)) + 1e-6
|
||||
centroid = float(np.sum(freq_values * weights) / denom)
|
||||
brightness = float(np.clip((centroid - 500.0) / 2500.0, 0.0, 1.0))
|
||||
return energy, brightness
|
||||
|
||||
def _controls_from_features(self, energy: float, brightness: float, phase: float) -> dict[str, float]:
|
||||
jaw_open = float(np.clip(energy, 0.0, 1.0))
|
||||
mouth_close = float(np.clip(1.0 - jaw_open * 0.82, 0.0, 1.0))
|
||||
mouth_funnel = float(np.clip(0.16 + jaw_open * (0.40 - brightness * 0.12), 0.0, 1.0))
|
||||
mouth_pucker = float(np.clip(0.06 + jaw_open * (0.20 + (1.0 - brightness) * 0.30), 0.0, 1.0))
|
||||
viseme_aa = jaw_open
|
||||
viseme_ee = float(np.clip(jaw_open * (0.25 + brightness * 0.95), 0.0, 1.0))
|
||||
viseme_oh = float(np.clip(jaw_open * (0.30 + (1.0 - brightness) * 0.90), 0.0, 1.0))
|
||||
head_yaw = float(np.clip(math.sin(phase * math.pi * 1.7) * 0.02 * max(jaw_open, 0.18), -0.25, 0.25))
|
||||
head_pitch = float(np.clip(math.cos(phase * math.pi * 2.0) * 0.03 * max(jaw_open, 0.15), -0.25, 0.25))
|
||||
head_roll = float(np.clip(math.sin(phase * math.pi) * 0.015, -0.15, 0.15))
|
||||
return {
|
||||
"jawOpen": round(jaw_open, 4),
|
||||
"mouthClose": round(mouth_close, 4),
|
||||
"mouthFunnel": round(mouth_funnel, 4),
|
||||
"mouthPucker": round(mouth_pucker, 4),
|
||||
"viseme_aa": round(viseme_aa, 4),
|
||||
"viseme_ee": round(viseme_ee, 4),
|
||||
"viseme_oh": round(viseme_oh, 4),
|
||||
"headYaw": round(head_yaw, 4),
|
||||
"headPitch": round(head_pitch, 4),
|
||||
"headRoll": round(head_roll, 4),
|
||||
}
|
||||
|
||||
def _build_frame(self, seq: int, time_ms: int, controls: dict[str, float]) -> dict[str, Any]:
|
||||
return {
|
||||
"seq": seq,
|
||||
"time_ms": time_ms,
|
||||
"speaking": controls["jawOpen"] > 0.04,
|
||||
"controls": controls,
|
||||
}
|
||||
|
||||
def _neutral_controls(self) -> dict[str, float]:
|
||||
return {
|
||||
"jawOpen": 0.0,
|
||||
"mouthClose": 1.0,
|
||||
"mouthFunnel": 0.0,
|
||||
"mouthPucker": 0.0,
|
||||
"viseme_aa": 0.0,
|
||||
"viseme_ee": 0.0,
|
||||
"viseme_oh": 0.0,
|
||||
"headYaw": 0.0,
|
||||
"headPitch": 0.0,
|
||||
"headRoll": 0.0,
|
||||
}
|
||||
|
||||
async def _next_sequence(self) -> int:
|
||||
async with self._lock:
|
||||
seq = self._sequence
|
||||
self._sequence += 1
|
||||
return seq
|
||||
|
||||
async def _reserve_sequences(self, count: int) -> range:
|
||||
async with self._lock:
|
||||
start = self._sequence
|
||||
self._sequence += count
|
||||
return range(start, start + count)
|
||||
|
||||
@property
|
||||
def health(self) -> dict[str, Any]:
|
||||
return {
|
||||
"driver": self.driver_name,
|
||||
"mode": settings.avatar_driver_mode,
|
||||
"protocol": settings.avatar_control_protocol,
|
||||
"schema": settings.avatar_blendshape_schema,
|
||||
"last_frame_count": self._last_frame_count,
|
||||
"last_chunk_duration_ms": self._last_chunk_duration_ms,
|
||||
"last_error": self._last_error,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user