#!/usr/bin/env python3 """ This is an extra gRPC server of LocalAI for Qwen3-TTS """ 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 import soundfile as sf from qwen_tts import Qwen3TTSModel 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 # Get model path from request model_path = request.Model if not model_path: model_path = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" # Determine model type from model path or options self.model_type = self.options.get("model_type", None) if not self.model_type: if "CustomVoice" in model_path: self.model_type = "CustomVoice" elif "VoiceDesign" in model_path: self.model_type = "VoiceDesign" elif "Base" in model_path or "0.6B" in model_path or "1.7B" in model_path: self.model_type = "Base" # VoiceClone model else: # Default to CustomVoice self.model_type = "CustomVoice" # Cache for voice clone prompts self._voice_clone_cache = {} # Store AudioPath, ModelFile, and ModelPath from LoadModel request # These are used later in TTS for VoiceClone mode 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 dtype & attention implementation if self.device == "mps": load_dtype = torch.float32 # MPS requires float32 device_map = None attn_impl_primary = "sdpa" # flash_attention_2 not supported on MPS elif self.device == "cuda": load_dtype = torch.bfloat16 device_map = "cuda" attn_impl_primary = "flash_attention_2" else: # cpu load_dtype = torch.float32 device_map = "cpu" attn_impl_primary = "sdpa" print(f"Using device: {self.device}, torch_dtype: {load_dtype}, attn_implementation: {attn_impl_primary}, model_type: {self.model_type}", file=sys.stderr) print(f"Loading model from: {model_path}", file=sys.stderr) # Load model with device-specific logic # Common parameters for all devices load_kwargs = { "dtype": load_dtype, "attn_implementation": attn_impl_primary, "trust_remote_code": True, # Required for qwen-tts models } try: if self.device == "mps": load_kwargs["device_map"] = None # load then move self.model = Qwen3TTSModel.from_pretrained(model_path, **load_kwargs) self.model.to("mps") elif self.device == "cuda": load_kwargs["device_map"] = device_map self.model = Qwen3TTSModel.from_pretrained(model_path, **load_kwargs) else: # cpu load_kwargs["device_map"] = device_map self.model = Qwen3TTSModel.from_pretrained(model_path, **load_kwargs) except Exception as e: error_msg = str(e) print(f"[ERROR] Loading model: {type(e).__name__}: {error_msg}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) # Check if it's a missing feature extractor/tokenizer error if "speech_tokenizer" in error_msg or "preprocessor_config.json" in error_msg or "feature extractor" in error_msg.lower(): print("\n[ERROR] Model files appear to be incomplete. This usually means:", file=sys.stderr) print(" 1. The model download was interrupted or incomplete", file=sys.stderr) print(" 2. The model cache is corrupted", file=sys.stderr) print("\nTo fix this, try:", file=sys.stderr) print(f" rm -rf ~/.cache/huggingface/hub/models--Qwen--Qwen3-TTS-*", file=sys.stderr) print(" Then re-run to trigger a fresh download.", file=sys.stderr) print("\nAlternatively, try using a different model variant:", file=sys.stderr) print(" - Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", file=sys.stderr) print(" - Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign", file=sys.stderr) print(" - Qwen/Qwen3-TTS-12Hz-1.7B-Base", file=sys.stderr) if attn_impl_primary == 'flash_attention_2': print("\nTrying to use SDPA instead of flash_attention_2...", file=sys.stderr) load_kwargs["attn_implementation"] = 'sdpa' try: if self.device == "mps": load_kwargs["device_map"] = None self.model = Qwen3TTSModel.from_pretrained(model_path, **load_kwargs) self.model.to("mps") else: load_kwargs["device_map"] = (self.device if self.device in ("cuda", "cpu") else None) self.model = Qwen3TTSModel.from_pretrained(model_path, **load_kwargs) except Exception as e2: print(f"[ERROR] Failed to load with SDPA: {type(e2).__name__}: {e2}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) raise e2 else: raise e print(f"Model loaded successfully: {model_path}", file=sys.stderr) return backend_pb2.Result(message="Model loaded successfully", success=True) def _detect_mode(self, request): """Detect which mode to use based on request parameters.""" # Priority: VoiceClone > VoiceDesign > CustomVoice # model_type explicitly set if self.model_type == "CustomVoice": return "CustomVoice" if self.model_type == "VoiceClone": return "VoiceClone" if self.model_type == "VoiceDesign": return "VoiceDesign" # VoiceClone: AudioPath is provided (from LoadModel, stored in self.audio_path) if self.audio_path: return "VoiceClone" # VoiceDesign: instruct option is provided if "instruct" in self.options and self.options["instruct"]: return "VoiceDesign" # Default to CustomVoice return "CustomVoice" def _get_ref_audio_path(self, request): """Get reference audio path from stored AudioPath (from LoadModel).""" if not self.audio_path: return None # If absolute path, use as-is if os.path.isabs(self.audio_path): return self.audio_path # Try relative to ModelFile if self.model_file: model_file_base = os.path.dirname(self.model_file) ref_path = os.path.join(model_file_base, self.audio_path) if os.path.exists(ref_path): return ref_path # Try relative to ModelPath if self.model_path: ref_path = os.path.join(self.model_path, self.audio_path) if os.path.exists(ref_path): return ref_path # Return as-is (might be URL or base64) return self.audio_path def _get_voice_clone_prompt(self, request, ref_audio, ref_text): """Get or create voice clone prompt, with caching.""" cache_key = f"{ref_audio}:{ref_text}" if cache_key not in self._voice_clone_cache: print(f"Creating voice clone prompt from {ref_audio}", file=sys.stderr) try: prompt_items = self.model.create_voice_clone_prompt( ref_audio=ref_audio, ref_text=ref_text, x_vector_only_mode=self.options.get("x_vector_only_mode", False), ) self._voice_clone_cache[cache_key] = prompt_items except Exception as e: print(f"Error creating voice clone prompt: {e}", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) return None return self._voice_clone_cache[cache_key] def TTS(self, request, context): try: # Check if dst is provided if not request.dst: return backend_pb2.Result( success=False, message="dst (output path) is required" ) # Prepare text text = request.text.strip() if not text: return backend_pb2.Result( success=False, message="Text is empty" ) # Get language (auto-detect if not provided) language = request.language if hasattr(request, 'language') and request.language else None if not language or language == "": language = "Auto" # Auto-detect language # Detect mode mode = self._detect_mode(request) print(f"Detected mode: {mode}", file=sys.stderr) # Get generation parameters from options max_new_tokens = self.options.get("max_new_tokens", None) top_p = self.options.get("top_p", None) temperature = self.options.get("temperature", None) do_sample = self.options.get("do_sample", None) # Prepare generation kwargs generation_kwargs = {} if max_new_tokens is not None: generation_kwargs["max_new_tokens"] = max_new_tokens if top_p is not None: generation_kwargs["top_p"] = top_p if temperature is not None: generation_kwargs["temperature"] = temperature if do_sample is not None: generation_kwargs["do_sample"] = do_sample instruct = self.options.get("instruct", "") if instruct is not None and instruct != "": generation_kwargs["instruct"] = instruct # Generate audio based on mode if mode == "VoiceClone": # VoiceClone mode ref_audio = self._get_ref_audio_path(request) if not ref_audio: return backend_pb2.Result( success=False, message="AudioPath is required for VoiceClone mode" ) ref_text = self.options.get("ref_text", None) if not ref_text: # Try to get from request if available if hasattr(request, 'ref_text') and request.ref_text: ref_text = request.ref_text else: # x_vector_only_mode doesn't require ref_text if not self.options.get("x_vector_only_mode", False): return backend_pb2.Result( success=False, message="ref_text is required for VoiceClone mode (or set x_vector_only_mode=true)" ) # Check if we should use cached prompt use_cached_prompt = self.options.get("use_cached_prompt", True) voice_clone_prompt = None if use_cached_prompt: voice_clone_prompt = self._get_voice_clone_prompt(request, ref_audio, ref_text) if voice_clone_prompt is None: return backend_pb2.Result( success=False, message="Failed to create voice clone prompt" ) if voice_clone_prompt: # Use cached prompt wavs, sr = self.model.generate_voice_clone( text=text, language=language, voice_clone_prompt=voice_clone_prompt, **generation_kwargs ) else: # Create prompt on-the-fly wavs, sr = self.model.generate_voice_clone( text=text, language=language, ref_audio=ref_audio, ref_text=ref_text, x_vector_only_mode=self.options.get("x_vector_only_mode", False), **generation_kwargs ) elif mode == "VoiceDesign": # VoiceDesign mode if not instruct: return backend_pb2.Result( success=False, message="instruct option is required for VoiceDesign mode" ) wavs, sr = self.model.generate_voice_design( text=text, language=language, instruct=instruct, **generation_kwargs ) else: # CustomVoice mode (default) speaker = request.voice if request.voice else None if not speaker: # Try to get from options speaker = self.options.get("speaker", None) if not speaker: # Use default speaker speaker = "Vivian" print(f"No speaker specified, using default: {speaker}", file=sys.stderr) # Validate speaker if model supports it if hasattr(self.model, 'get_supported_speakers'): try: supported_speakers = self.model.get_supported_speakers() if speaker not in supported_speakers: print(f"Warning: Speaker '{speaker}' not in supported list. Available: {supported_speakers}", file=sys.stderr) # Try to find a close match (case-insensitive) speaker_lower = speaker.lower() for sup_speaker in supported_speakers: if sup_speaker.lower() == speaker_lower: speaker = sup_speaker print(f"Using matched speaker: {speaker}", file=sys.stderr) break except Exception as e: print(f"Warning: Could not get supported speakers: {e}", file=sys.stderr) wavs, sr = self.model.generate_custom_voice( text=text, language=language, speaker=speaker, **generation_kwargs ) # Save output if wavs is not None and len(wavs) > 0: # wavs is a list, take first element audio_data = wavs[0] if isinstance(wavs, list) else wavs sf.write(request.dst, audio_data, sr) 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 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)