#!/usr/bin/env python3 """ Single-user streaming test for MuseTalk gRPC Avatar Service Uses DIRECT gRPC connection (no SSH tunnel required) Saves output video with audio to timestamped directory Usage: python benchmarks/avatar_single_test.py --server 81.166.173.12:10597 --audio data/audio/test.wav """ import sys import os import time import wave import asyncio import argparse import shutil import subprocess from datetime import datetime from io import BytesIO # Add grpc module path sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'server', 'grpc')) import grpc from grpc import aio import avatar_pb2 import avatar_pb2_grpc def create_output_dir(base_dir: str, test_name: str) -> str: """Create timestamped output directory""" timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") output_dir = os.path.join(base_dir, f"{timestamp}_{test_name}") os.makedirs(output_dir, exist_ok=True) return output_dir def save_video_with_audio(frames: list, audio_path: str, output_dir: str, fps: int = 25): """Save frames as video and combine with audio using ffmpeg""" try: import cv2 import numpy as np except ImportError: print("[Warning] opencv-python not installed, skipping video save") return None if not frames: print("[Warning] No frames to save") return None # Decode first frame to get dimensions first_frame = cv2.imdecode(np.frombuffer(frames[0], np.uint8), cv2.IMREAD_COLOR) if first_frame is None: print("[Warning] Failed to decode first frame") return None height, width = first_frame.shape[:2] # Save frames as temporary video (no audio) temp_video = os.path.join(output_dir, "temp_video.mp4") fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(temp_video, fourcc, fps, (width, height)) for i, frame_data in enumerate(frames): frame = cv2.imdecode(np.frombuffer(frame_data, np.uint8), cv2.IMREAD_COLOR) if frame is not None: out.write(frame) out.release() # Combine video with audio using ffmpeg output_video = os.path.join(output_dir, "output.mp4") try: cmd = [ 'ffmpeg', '-y', '-i', temp_video, '-i', audio_path, '-c:v', 'libx264', '-c:a', 'aac', '-shortest', '-preset', 'fast', output_video ] subprocess.run(cmd, capture_output=True, check=True) os.remove(temp_video) return output_video except subprocess.CalledProcessError as e: print(f"[Warning] ffmpeg failed: {e}") # Keep temp video if ffmpeg fails return temp_video except FileNotFoundError: print("[Warning] ffmpeg not found, keeping video without audio") shutil.move(temp_video, output_video) return output_video async def main(): parser = argparse.ArgumentParser(description='MuseTalk Avatar Single-User Test') parser.add_argument('--server', type=str, default='81.166.173.12:10597', help='gRPC server address (default: 81.166.173.12:10597)') parser.add_argument('--audio', type=str, default=os.path.join(os.path.dirname(__file__), '..', 'orpheus_demo_3_phrases.wav'), help='Path to audio file for testing') parser.add_argument('--chunk-ms', type=int, default=100, help='Audio chunk size in milliseconds (default: 100)') parser.add_argument('--output-dir', type=str, default=os.path.join(os.path.dirname(__file__), 'results'), help='Base directory for output (default: benchmarks/results)') parser.add_argument('--no-save', action='store_true', help='Skip saving video output') args = parser.parse_args() print("=" * 60) print("MuseTalk gRPC Single-User Test") print(f"Server: {args.server}") print("=" * 60) # Create output directory if not args.no_save: output_dir = create_output_dir(args.output_dir, "single_user") print(f"[Output] Saving to: {output_dir}") # Load audio audio_path = os.path.abspath(args.audio) with wave.open(audio_path, 'rb') as wf: sample_rate = wf.getframerate() n_frames = wf.getnframes() duration = n_frames / sample_rate audio_data = wf.readframes(n_frames) print(f"[Audio] Loaded {args.audio}") print(f" Sample rate: {sample_rate} Hz") print(f" Duration: {duration:.2f}s") print(f" Samples: {n_frames}") # Copy audio to output dir if not args.no_save: audio_copy = os.path.join(output_dir, os.path.basename(audio_path)) shutil.copy2(audio_path, audio_copy) # Connect to server print(f"\n[gRPC] Connecting to {args.server}...") channel = aio.insecure_channel(args.server) stub = avatar_pb2_grpc.AvatarServiceStub(channel) # Health check try: health = await stub.HealthCheck(avatar_pb2.HealthRequest()) print(f"[gRPC] Server status: {health.status}, avatar: {health.avatar_id}") except Exception as e: print(f"[gRPC] Health check failed: {e}") return # Get avatar info info = await stub.GetAvatarInfo(avatar_pb2.AvatarInfoRequest()) print(f"[gRPC] Avatar: {info.avatar_id}, resolution: {info.width}x{info.height}, fps: {info.fps}") # Stream audio chunk_ms = args.chunk_ms chunk_size = int(sample_rate * chunk_ms / 1000) * 2 # bytes (16-bit) session_id = f"single_test_{int(time.time())}" async def audio_generator(): for i in range(0, len(audio_data), chunk_size): chunk = audio_data[i:i + chunk_size] is_final = (i + chunk_size >= len(audio_data)) yield avatar_pb2.AudioChunk( session_id=session_id, audio_data=chunk, sample_rate=sample_rate, is_final=is_final, timestamp=i / 2 / sample_rate, chunk_index=i // chunk_size ) # Simulate real-time streaming (slightly faster) await asyncio.sleep(chunk_ms / 1000 * 0.5) print(f"\n[Stream] Starting audio stream...") start_time = time.time() first_frame_time = None frame_count = 0 total_bytes = 0 frames = [] # Store frames for video try: async for response in stub.StreamingGenerate(audio_generator()): if first_frame_time is None: first_frame_time = time.time() ttff = (first_frame_time - start_time) * 1000 print(f"[Stream] Time to first frame: {ttff:.0f}ms") frame_count += 1 total_bytes += len(response.frame_data) # Store frame data for video if not args.no_save: frames.append(response.frame_data) if frame_count % 50 == 0: print(f"[Stream] Received frame {frame_count}...") if response.is_final: break except Exception as e: print(f"[Stream] Error: {e}") return elapsed = time.time() - first_frame_time if first_frame_time else 0 print(f"\n[Stream] Complete!") print(f" Total frames: {frame_count}") print(f" Elapsed time: {elapsed:.2f}s") print(f" Effective FPS: {frame_count/elapsed:.1f}" if elapsed > 0 else "N/A") print(f" Total data: {total_bytes/1024/1024:.2f} MB") print(f" Real-time ratio: {elapsed/duration:.2f}x") # Save video with audio if not args.no_save and frames: print(f"\n[Video] Saving {len(frames)} frames...") video_path = save_video_with_audio(frames, audio_path, output_dir, fps=info.fps) if video_path: print(f"[Video] Saved to: {video_path}") # Save metrics metrics_path = os.path.join(output_dir, "metrics.txt") with open(metrics_path, 'w') as f: f.write(f"Test: Single User Streaming\n") f.write(f"Date: {datetime.now().isoformat()}\n") f.write(f"Server: {args.server}\n") f.write(f"Avatar: {info.avatar_id}\n") f.write(f"Resolution: {info.width}x{info.height}\n") f.write(f"Target FPS: {info.fps}\n") f.write(f"\n--- Audio ---\n") f.write(f"File: {os.path.basename(audio_path)}\n") f.write(f"Duration: {duration:.2f}s\n") f.write(f"Sample rate: {sample_rate} Hz\n") f.write(f"\n--- Results ---\n") f.write(f"TTFF: {ttff:.0f}ms\n") f.write(f"Total frames: {frame_count}\n") f.write(f"Elapsed time: {elapsed:.2f}s\n") f.write(f"Effective FPS: {frame_count/elapsed:.1f}\n") f.write(f"Total data: {total_bytes/1024/1024:.2f} MB\n") f.write(f"Real-time ratio: {elapsed/duration:.2f}x\n") print(f"[Metrics] Saved to: {metrics_path}") await channel.close() if __name__ == '__main__': asyncio.run(main())