#!/usr/bin/env python3 """ This is an extra gRPC server of LocalAI for VibeVoice """ from concurrent import futures import time import argparse import signal import sys import os import copy import traceback from pathlib import Path import backend_pb2 import backend_pb2_grpc import torch from vibevoice.modular.modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference from vibevoice.processor.vibevoice_streaming_processor import VibeVoiceStreamingProcessor from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor import grpc def is_float(s): """Check if a string can be converted to float.""" try: float(s) return True except ValueError: return False def is_int(s): """Check if a string can be converted to int.""" try: int(s) return True except ValueError: return False _ONE_DAY_IN_SECONDS = 60 * 60 * 24 # If MAX_WORKERS are specified in the environment use it, otherwise default to 1 MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1')) # Implement the BackendServicer class with the service methods class BackendServicer(backend_pb2_grpc.BackendServicer): """ BackendServicer is the class that implements the gRPC service """ def Health(self, request, context): return backend_pb2.Reply(message=bytes("OK", 'utf-8')) def LoadModel(self, request, context): # Get device if torch.cuda.is_available(): print("CUDA is available", file=sys.stderr) device = "cuda" else: print("CUDA is not available", file=sys.stderr) device = "cpu" mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available() if mps_available: device = "mps" if not torch.cuda.is_available() and request.CUDA: return backend_pb2.Result(success=False, message="CUDA is not available") # Normalize potential 'mpx' typo to 'mps' if device == "mpx": print("Note: device 'mpx' detected, treating it as 'mps'.", file=sys.stderr) device = "mps" # Validate mps availability if requested if device == "mps" and not torch.backends.mps.is_available(): print("Warning: MPS not available. Falling back to CPU.", file=sys.stderr) device = "cpu" self.device = device self._torch_device = torch.device(device) options = request.Options # empty dict self.options = {} # The options are a list of strings in this form optname:optvalue # We are storing all the options in a dict so we can use it later when # generating the audio for opt in options: if ":" not in opt: continue key, value = opt.split(":", 1) # Split only on first colon # if value is a number, convert it to the appropriate type if is_float(value): value = float(value) elif is_int(value): value = int(value) elif value.lower() in ["true", "false"]: value = value.lower() == "true" self.options[key] = value # Check if ASR mode is enabled self.asr_mode = self.options.get("asr_mode", False) if not isinstance(self.asr_mode, bool): # Handle string "true"/"false" case self.asr_mode = str(self.asr_mode).lower() == "true" # Get model path from request model_path = request.Model if not model_path: if self.asr_mode: model_path = "microsoft/VibeVoice-ASR" # Default ASR model else: model_path = "microsoft/VibeVoice-Realtime-0.5B" # Default TTS model default_dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 load_dtype = default_dtype if "torch_dtype" in self.options: torch_dtype_str = str(self.options["torch_dtype"]).lower() if torch_dtype_str == "fp16": load_dtype = torch.float16 elif torch_dtype_str == "bf16": load_dtype = torch.bfloat16 elif torch_dtype_str == "fp32": load_dtype = torch.float32 # remove it from options after reading del self.options["torch_dtype"] # Get inference steps from options, default to 5 (TTS only) self.inference_steps = self.options.get("inference_steps", 5) if not isinstance(self.inference_steps, int) or self.inference_steps <= 0: self.inference_steps = 5 # Get cfg_scale from options, default to 1.5 (TTS only) self.cfg_scale = self.options.get("cfg_scale", 1.5) if not isinstance(self.cfg_scale, (int, float)) or self.cfg_scale <= 0: self.cfg_scale = 1.5 # Get ASR generation parameters from options self.max_new_tokens = self.options.get("max_new_tokens", 512) if not isinstance(self.max_new_tokens, int) or self.max_new_tokens <= 0: self.max_new_tokens = 512 self.temperature = self.options.get("temperature", 0.0) if not isinstance(self.temperature, (int, float)) or self.temperature < 0: self.temperature = 0.0 self.top_p = self.options.get("top_p", 1.0) if not isinstance(self.top_p, (int, float)) or self.top_p <= 0: self.top_p = 1.0 self.do_sample = self.options.get("do_sample", None) if self.do_sample is None: # Default: use sampling if temperature > 0 self.do_sample = self.temperature > 0 elif not isinstance(self.do_sample, bool): self.do_sample = str(self.do_sample).lower() == "true" self.num_beams = self.options.get("num_beams", 1) if not isinstance(self.num_beams, int) or self.num_beams < 1: self.num_beams = 1 self.repetition_penalty = self.options.get("repetition_penalty", 1.0) if not isinstance(self.repetition_penalty, (int, float)) or self.repetition_penalty <= 0: self.repetition_penalty = 1.0 # Determine voices directory # Priority order: # 1. voices_dir option (explicitly set by user - highest priority) # 2. Relative to ModelFile if provided # 3. Relative to ModelPath (models directory) if provided # 4. Backend directory # 5. Absolute path from AudioPath if provided voices_dir = None # First check if voices_dir is explicitly set in options if "voices_dir" in self.options: voices_dir_option = self.options["voices_dir"] if isinstance(voices_dir_option, str) and voices_dir_option.strip(): voices_dir = voices_dir_option.strip() # If relative path, try to resolve it relative to ModelPath or ModelFile if not os.path.isabs(voices_dir): if hasattr(request, 'ModelPath') and request.ModelPath: voices_dir = os.path.join(request.ModelPath, voices_dir) elif request.ModelFile: model_file_base = os.path.dirname(request.ModelFile) voices_dir = os.path.join(model_file_base, voices_dir) # If still relative, make it absolute from current working directory if not os.path.isabs(voices_dir): voices_dir = os.path.abspath(voices_dir) # Check if the directory exists if not os.path.exists(voices_dir): print(f"Warning: voices_dir option specified but directory does not exist: {voices_dir}", file=sys.stderr) voices_dir = None # If not set via option, try relative to ModelFile if provided if not voices_dir and request.ModelFile: model_file_base = os.path.dirname(request.ModelFile) voices_dir = os.path.join(model_file_base, "voices", "streaming_model") if not os.path.exists(voices_dir): voices_dir = None # If not found, try relative to ModelPath (models directory) if not voices_dir and hasattr(request, 'ModelPath') and request.ModelPath: voices_dir = os.path.join(request.ModelPath, "voices", "streaming_model") if not os.path.exists(voices_dir): voices_dir = None # If not found, try relative to backend directory if not voices_dir: backend_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) voices_dir = os.path.join(backend_dir, "vibevoice", "voices", "streaming_model") if not os.path.exists(voices_dir): # Try absolute path from AudioPath if provided if request.AudioPath and os.path.isabs(request.AudioPath): voices_dir = os.path.dirname(request.AudioPath) else: voices_dir = None # Initialize voice-related attributes (TTS only) self.voices_dir = voices_dir self.voice_presets = {} self._voice_cache = {} self.default_voice_key = None # Store AudioPath, ModelFile, and ModelPath from LoadModel request for use in TTS self.audio_path = request.AudioPath if hasattr(request, 'AudioPath') and request.AudioPath else None self.model_file = request.ModelFile if hasattr(request, 'ModelFile') and request.ModelFile else None self.model_path = request.ModelPath if hasattr(request, 'ModelPath') and request.ModelPath else None # Decide attention implementation and device_map (matching upstream example) if self.device == "mps": device_map = None attn_impl_primary = "sdpa" # flash_attention_2 not supported on MPS elif self.device == "cuda": device_map = "cuda" attn_impl_primary = "flash_attention_2" else: # cpu device_map = "cpu" # Match upstream example: use "cpu" for CPU device_map attn_impl_primary = "sdpa" try: if self.asr_mode: # Load ASR model and processor print(f"Loading ASR processor & model from {model_path}", file=sys.stderr) # Load ASR processor self.processor = VibeVoiceASRProcessor.from_pretrained( model_path, language_model_pretrained_name="Qwen/Qwen2.5-7B" ) print(f"Using device: {self.device}, torch_dtype: {load_dtype}, attn_implementation: {attn_impl_primary}", file=sys.stderr) # Load ASR model - use device_map=None and move manually to avoid JSON serialization issues # Load with dtype to ensure all components are in correct dtype from the start try: print(f"Using attention implementation: {attn_impl_primary}", file=sys.stderr) # Load model with dtype to ensure all components are in correct dtype self.model = VibeVoiceASRForConditionalGeneration.from_pretrained( model_path, dtype=load_dtype, device_map=None, # Always use None, move manually to avoid JSON serialization issues attn_implementation=attn_impl_primary, trust_remote_code=True ) # Move to device manually self.model = self.model.to(self.device) except Exception as e: if attn_impl_primary == 'flash_attention_2': print(f"[ERROR] : {type(e).__name__}: {e}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) print("Error loading the ASR model. Trying to use SDPA.", file=sys.stderr) self.model = VibeVoiceASRForConditionalGeneration.from_pretrained( model_path, dtype=load_dtype, device_map=None, attn_implementation='sdpa', trust_remote_code=True ) # Move to device manually self.model = self.model.to(self.device) else: raise e self.model.eval() print(f"ASR model loaded successfully", file=sys.stderr) else: # Load TTS model and processor (existing logic) # Load voice presets if directory exists if self.voices_dir and os.path.exists(self.voices_dir): self._load_voice_presets() else: print(f"Warning: Voices directory not found. Voice presets will not be available.", file=sys.stderr) print(f"Loading TTS processor & model from {model_path}", file=sys.stderr) self.processor = VibeVoiceStreamingProcessor.from_pretrained(model_path) print(f"Using device: {self.device}, torch_dtype: {load_dtype}, attn_implementation: {attn_impl_primary}", file=sys.stderr) # Load model with device-specific logic (matching upstream example exactly) try: if self.device == "mps": self.model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( model_path, torch_dtype=load_dtype, attn_implementation=attn_impl_primary, device_map=None, # load then move ) self.model.to("mps") elif self.device == "cuda": self.model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( model_path, torch_dtype=load_dtype, device_map=device_map, attn_implementation=attn_impl_primary, ) else: # cpu # Match upstream example: use device_map="cpu" for CPU self.model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( model_path, torch_dtype=load_dtype, device_map="cpu", attn_implementation=attn_impl_primary, ) except Exception as e: if attn_impl_primary == 'flash_attention_2': print(f"[ERROR] : {type(e).__name__}: {e}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) print("Error loading the model. Trying to use SDPA. However, note that only flash_attention_2 has been fully tested, and using SDPA may result in lower audio quality.", file=sys.stderr) # Match upstream example fallback pattern self.model = VibeVoiceStreamingForConditionalGenerationInference.from_pretrained( model_path, torch_dtype=load_dtype, device_map=(self.device if self.device in ("cuda", "cpu") else None), attn_implementation='sdpa' ) if self.device == "mps": self.model.to("mps") else: raise e self.model.eval() self.model.set_ddpm_inference_steps(num_steps=self.inference_steps) # Set default voice key if self.voice_presets: # Try to get default from environment or use first available preset_name = os.environ.get("VOICE_PRESET") self.default_voice_key = self._determine_voice_key(preset_name) print(f"Default voice preset: {self.default_voice_key}", file=sys.stderr) else: print("Warning: No voice presets available. Voice selection will not work.", file=sys.stderr) except Exception as err: # Format error message safely, avoiding JSON serialization issues error_msg = str(err) error_type = type(err).__name__ # Include traceback for debugging tb_str = traceback.format_exc() print(f"[ERROR] LoadModel failed: {error_type}: {error_msg}", file=sys.stderr) print(tb_str, file=sys.stderr) return backend_pb2.Result(success=False, message=f"{error_type}: {error_msg}") return backend_pb2.Result(message="Model loaded successfully", success=True) def _load_voice_presets(self): """Load voice presets from the voices directory.""" if not self.voices_dir or not os.path.exists(self.voices_dir): self.voice_presets = {} return self.voice_presets = {} # Get all .pt files in the voices directory pt_files = [f for f in os.listdir(self.voices_dir) if f.lower().endswith('.pt') and os.path.isfile(os.path.join(self.voices_dir, f))] # Create dictionary with filename (without extension) as key for pt_file in pt_files: # Remove .pt extension to get the name name = os.path.splitext(pt_file)[0] # Create full path full_path = os.path.join(self.voices_dir, pt_file) self.voice_presets[name] = full_path # Sort the voice presets alphabetically by name self.voice_presets = dict(sorted(self.voice_presets.items())) print(f"Found {len(self.voice_presets)} voice files in {self.voices_dir}", file=sys.stderr) if self.voice_presets: print(f"Available voices: {', '.join(self.voice_presets.keys())}", file=sys.stderr) def _determine_voice_key(self, name): """Determine voice key from name or use default.""" if name and name in self.voice_presets: return name # Try default key default_key = "en-WHTest_man" if default_key in self.voice_presets: return default_key # Use first available if self.voice_presets: first_key = next(iter(self.voice_presets)) print(f"Using fallback voice preset: {first_key}", file=sys.stderr) return first_key return None def _get_voice_path(self, speaker_name): """Get voice file path for a given speaker name.""" if not self.voice_presets: return None # First try exact match if speaker_name and speaker_name in self.voice_presets: return self.voice_presets[speaker_name] # Try partial matching (case insensitive) if speaker_name: speaker_lower = speaker_name.lower() for preset_name, path in self.voice_presets.items(): if preset_name.lower() in speaker_lower or speaker_lower in preset_name.lower(): return path # Default to first voice if no match found if self.default_voice_key and self.default_voice_key in self.voice_presets: return self.voice_presets[self.default_voice_key] elif self.voice_presets: default_voice = list(self.voice_presets.values())[0] print(f"Warning: No voice preset found for '{speaker_name}', using default voice: {default_voice}", file=sys.stderr) return default_voice return None def _ensure_voice_cached(self, voice_path): """Load and cache voice preset.""" if not voice_path or not os.path.exists(voice_path): return None # Ensure cache exists (should be initialized in LoadModel) if not hasattr(self, '_voice_cache'): self._voice_cache = {} # Use path as cache key if voice_path not in self._voice_cache: print(f"Loading prefilled prompt from {voice_path}", file=sys.stderr) # Match self-test.py: use string device name for map_location # Ensure self.device exists (should be set in LoadModel) try: if not hasattr(self, 'device'): # Fallback to CPU if device not set device_str = "cpu" else: device_str = str(self.device) except AttributeError as e: print(f"Error accessing self.device: {e}, falling back to CPU", file=sys.stderr) device_str = "cpu" if device_str != "cpu": map_loc = device_str else: map_loc = "cpu" # Call torch.load with explicit arguments prefilled_outputs = torch.load(voice_path, map_location=map_loc, weights_only=False) self._voice_cache[voice_path] = prefilled_outputs return self._voice_cache[voice_path] def TTS(self, request, context): try: # Get voice selection # Priority: request.voice > AudioPath > default voice_path = None voice_key = None if request.voice: # Try to get voice by name voice_path = self._get_voice_path(request.voice) if voice_path: voice_key = request.voice elif self.audio_path: # Use AudioPath from LoadModel as voice file if os.path.isabs(self.audio_path): voice_path = self.audio_path elif self.model_file: model_file_base = os.path.dirname(self.model_file) voice_path = os.path.join(model_file_base, self.audio_path) elif self.model_path: voice_path = os.path.join(self.model_path, self.audio_path) else: voice_path = self.audio_path elif self.default_voice_key: voice_path = self._get_voice_path(self.default_voice_key) voice_key = self.default_voice_key if not voice_path or not os.path.exists(voice_path): return backend_pb2.Result( success=False, message=f"Voice file not found: {voice_path}. Please provide a valid voice preset or AudioPath." ) # Load voice preset prefilled_outputs = self._ensure_voice_cached(voice_path) if prefilled_outputs is None: return backend_pb2.Result( success=False, message=f"Failed to load voice preset from {voice_path}" ) # Get generation parameters from options cfg_scale = self.options.get("cfg_scale", self.cfg_scale) inference_steps = self.options.get("inference_steps", self.inference_steps) do_sample = self.options.get("do_sample", False) temperature = self.options.get("temperature", 0.9) top_p = self.options.get("top_p", 0.9) # Update inference steps if needed if inference_steps != self.inference_steps: self.model.set_ddpm_inference_steps(num_steps=inference_steps) self.inference_steps = inference_steps # Prepare text text = request.text.strip().replace("'", "'").replace('"', '"').replace('"', '"') # Prepare inputs inputs = self.processor.process_input_with_cached_prompt( text=text, cached_prompt=prefilled_outputs, padding=True, return_tensors="pt", return_attention_mask=True, ) # Move tensors to target device (matching self-test.py exactly) # Explicitly ensure it's a string to avoid any variable name collisions target_device = str(self.device) if str(self.device) != "cpu" else "cpu" for k, v in inputs.items(): if torch.is_tensor(v): inputs[k] = v.to(target_device) print(f"Generating audio with cfg_scale: {cfg_scale}, inference_steps: {inference_steps}", file=sys.stderr) # Generate audio outputs = self.model.generate( **inputs, max_new_tokens=None, cfg_scale=cfg_scale, tokenizer=self.processor.tokenizer, generation_config={ 'do_sample': do_sample, 'temperature': temperature if do_sample else 1.0, 'top_p': top_p if do_sample else 1.0, }, verbose=False, all_prefilled_outputs=copy.deepcopy(prefilled_outputs) if prefilled_outputs is not None else None, ) # Save output if outputs.speech_outputs and outputs.speech_outputs[0] is not None: self.processor.save_audio( outputs.speech_outputs[0], # First (and only) batch item output_path=request.dst, ) print(f"Saved output to {request.dst}", file=sys.stderr) else: return backend_pb2.Result( success=False, message="No audio output generated" ) except Exception as err: print(f"Error in TTS: {err}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}") return backend_pb2.Result(success=True) def AudioTranscription(self, request, context): """Transcribe audio file to text using ASR model.""" try: # Validate ASR mode is active if not self.asr_mode: return backend_pb2.TranscriptResult( segments=[], text="", ) # Note: We return empty result instead of error to match faster-whisper behavior # Get audio file path audio_path = request.dst if not audio_path or not os.path.exists(audio_path): print(f"Error: Audio file not found: {audio_path}", file=sys.stderr) return backend_pb2.TranscriptResult( segments=[], text="", ) print(f"Transcribing audio file: {audio_path}", file=sys.stderr) # Get context_info from options if available context_info = self.options.get("context_info", None) if context_info and isinstance(context_info, str) and context_info.strip(): context_info = context_info.strip() else: context_info = None # Process audio with ASR processor (matching gradio example) inputs = self.processor( audio=audio_path, sampling_rate=None, return_tensors="pt", add_generation_prompt=True, context_info=context_info ) # Move to device (matching gradio example) inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} # Prepare generation config (matching gradio example) generation_config = { "max_new_tokens": self.max_new_tokens, "temperature": self.temperature if self.temperature > 0 else None, "top_p": self.top_p if self.do_sample else None, "do_sample": self.do_sample, "num_beams": self.num_beams, "repetition_penalty": self.repetition_penalty, "pad_token_id": self.processor.pad_id, "eos_token_id": self.processor.tokenizer.eos_token_id, } # Remove None values (matching gradio example) generation_config = {k: v for k, v in generation_config.items() if v is not None} print(f"Generating transcription with max_new_tokens: {self.max_new_tokens}, temperature: {self.temperature}, do_sample: {self.do_sample}, num_beams: {self.num_beams}, repetition_penalty: {self.repetition_penalty}", file=sys.stderr) # Generate transcription (matching gradio example) with torch.no_grad(): output_ids = self.model.generate( **inputs, **generation_config ) # Decode output (matching gradio example) generated_ids = output_ids[0, inputs['input_ids'].shape[1]:] generated_text = self.processor.decode(generated_ids, skip_special_tokens=True) # Parse structured output to get segments result_segments = [] try: transcription_segments = self.processor.post_process_transcription(generated_text) if transcription_segments: # Map segments to TranscriptSegment format for idx, seg in enumerate(transcription_segments): # Extract timing information (if available) # Handle both dict and object with attributes if isinstance(seg, dict): start_time = seg.get('start_time', 0) end_time = seg.get('end_time', 0) text = seg.get('text', '') speaker_id = seg.get('speaker_id', None) else: # Handle object with attributes start_time = getattr(seg, 'start_time', 0) end_time = getattr(seg, 'end_time', 0) text = getattr(seg, 'text', '') speaker_id = getattr(seg, 'speaker_id', None) # Convert time to milliseconds (assuming seconds) start_ms = int(start_time * 1000) if isinstance(start_time, (int, float)) else 0 end_ms = int(end_time * 1000) if isinstance(end_time, (int, float)) else 0 # Add speaker info to text if available if speaker_id is not None: text = f"[Speaker {speaker_id}] {text}" result_segments.append(backend_pb2.TranscriptSegment( id=idx, start=start_ms, end=end_ms, text=text, tokens=[] # Token IDs not extracted for now )) except Exception as e: print(f"Warning: Failed to parse structured output: {e}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) # Fallback: create a single segment with the full text if generated_text: result_segments.append(backend_pb2.TranscriptSegment( id=0, start=0, end=0, text=generated_text, tokens=[] )) # Combine all segment texts into full transcription if result_segments: full_text = " ".join([seg.text for seg in result_segments]) else: full_text = generated_text if generated_text else "" print(f"Transcription completed: {len(result_segments)} segments", file=sys.stderr) return backend_pb2.TranscriptResult( segments=result_segments, text=full_text ) except Exception as err: print(f"Error in AudioTranscription: {err}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) return backend_pb2.TranscriptResult( segments=[], text="", ) def serve(address): server = grpc.server(futures.ThreadPoolExecutor(max_workers=MAX_WORKERS), options=[ ('grpc.max_message_length', 50 * 1024 * 1024), # 50MB ('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB ('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB ]) backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server) server.add_insecure_port(address) server.start() print("Server started. Listening on: " + address, file=sys.stderr) # Define the signal handler function def signal_handler(sig, frame): print("Received termination signal. Shutting down...") server.stop(0) sys.exit(0) # Set the signal handlers for SIGINT and SIGTERM signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) try: while True: time.sleep(_ONE_DAY_IN_SECONDS) except KeyboardInterrupt: server.stop(0) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Run the gRPC server.") parser.add_argument( "--addr", default="localhost:50051", help="The address to bind the server to." ) args = parser.parse_args() serve(args.addr)