LocalAI/backend/python/vllm-omni/backend.py

#!/usr/bin/env python3
"""
LocalAI vLLM-Omni Backend

This backend provides gRPC access to vllm-omni for multimodal generation:
- Image generation (text-to-image, image editing)
- Video generation (text-to-video, image-to-video)
- Text generation with multimodal inputs (LLM)
- Text-to-speech generation
"""
from concurrent import futures
import traceback
import argparse
import signal
import sys
import time
import os
import base64
import io
import json
import gc

from PIL import Image
import torch
import numpy as np
import soundfile as sf

import backend_pb2
import backend_pb2_grpc

import grpc
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
from grpc_auth import get_auth_interceptors
from vllm_utils import parse_options, messages_to_dicts, setup_parsers


from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.utils.platform_utils import detect_device_type, is_npu
from vllm import SamplingParams
from diffusers.utils import export_to_video

_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'))


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


# Implement the BackendServicer class with the service methods
class BackendServicer(backend_pb2_grpc.BackendServicer):

    def _detect_model_type(self, model_name):
        """Detect model type from model name."""
        model_lower = model_name.lower()
        if "tts" in model_lower or "qwen3-tts" in model_lower:
            return "tts"
        elif "omni" in model_lower and "qwen3" in model_lower:
            return "llm"
        elif "wan" in model_lower or "t2v" in model_lower or "i2v" in model_lower:
            return "video"
        elif "image" in model_lower or "z-image" in model_lower or "qwen-image" in model_lower:
            return "image"
        else:
            # Default to image for diffusion models, llm for others
            return "image"

    def _detect_tts_task_type(self):
        """Detect TTS task type from model name."""
        model_lower = self.model_name.lower()
        if "customvoice" in model_lower:
            return "CustomVoice"
        elif "voicedesign" in model_lower:
            return "VoiceDesign"
        elif "base" in model_lower:
            return "Base"
        else:
            # Default to CustomVoice
            return "CustomVoice"

    def _load_image(self, image_path):
        """Load an image from file path or base64 encoded data."""
        # Try file path first
        if os.path.exists(image_path):
            return Image.open(image_path)
        # Try base64 decode
        try:
            image_data = base64.b64decode(image_path)
            return Image.open(io.BytesIO(image_data))
        except:
            return None

    def _load_video(self, video_path):
        """Load a video from file path or base64 encoded data."""
        from vllm.assets.video import VideoAsset, video_to_ndarrays
        if os.path.exists(video_path):
            return video_to_ndarrays(video_path, num_frames=16)
        # Try base64 decode
        try:
            timestamp = str(int(time.time() * 1000))
            p = f"/tmp/vl-{timestamp}.data"
            with open(p, "wb") as f:
                f.write(base64.b64decode(video_path))
            video = VideoAsset(name=p).np_ndarrays
            os.remove(p)
            return video
        except:
            return None

    def _load_audio(self, audio_path):
        """Load audio from file path or base64 encoded data."""
        import librosa
        if os.path.exists(audio_path):
            audio_signal, sr = librosa.load(audio_path, sr=16000)
            return (audio_signal.astype(np.float32), sr)
        # Try base64 decode
        try:
            audio_data = base64.b64decode(audio_path)
            # Save to temp file and load
            timestamp = str(int(time.time() * 1000))
            p = f"/tmp/audio-{timestamp}.wav"
            with open(p, "wb") as f:
                f.write(audio_data)
            audio_signal, sr = librosa.load(p, sr=16000)
            os.remove(p)
            return (audio_signal.astype(np.float32), sr)
        except:
            return None

    def Health(self, request, context):
        return backend_pb2.Reply(message=bytes("OK", 'utf-8'))

    def LoadModel(self, request, context):
        try:
            # CPU detection: if no CUDA, default vLLM target device to CPU.
            try:
                if not torch.cuda.is_available():
                    os.environ.setdefault("VLLM_TARGET_DEVICE", "cpu")
                    os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "4")
            except Exception:
                pass

            print(f"Loading model {request.Model}...", file=sys.stderr)
            print(f"Request {request}", file=sys.stderr)

            # Parse options from request.Options using shared helper
            self.options = parse_options(request.Options)
            opts = self.options

            print(f"Options: {self.options}", file=sys.stderr)

            # Detect model type
            self.model_name = request.Model
            self.model_type = request.Type if request.Type else self._detect_model_type(request.Model)
            print(f"Detected model type: {self.model_type}", file=sys.stderr)

            # Build DiffusionParallelConfig if diffusion model (image or video)
            parallel_config = None
            if self.model_type in ["image", "video"]:
                parallel_config = DiffusionParallelConfig(
                    ulysses_degree=self.options.get("ulysses_degree", 1),
                    ring_degree=self.options.get("ring_degree", 1),
                    cfg_parallel_size=self.options.get("cfg_parallel_size", 1),
                    tensor_parallel_size=self.options.get("tensor_parallel_size", 1),
                )

            # Build cache_config dict if cache_backend specified
            cache_backend = self.options.get("cache_backend")  # "cache_dit" or "tea_cache"
            cache_config = None
            if cache_backend == "cache_dit":
                cache_config = {
                    "Fn_compute_blocks": self.options.get("cache_dit_fn_compute_blocks", 1),
                    "Bn_compute_blocks": self.options.get("cache_dit_bn_compute_blocks", 0),
                    "max_warmup_steps": self.options.get("cache_dit_max_warmup_steps", 4),
                    "residual_diff_threshold": self.options.get("cache_dit_residual_diff_threshold", 0.24),
                    "max_continuous_cached_steps": self.options.get("cache_dit_max_continuous_cached_steps", 3),
                    "enable_taylorseer": self.options.get("cache_dit_enable_taylorseer", False),
                    "taylorseer_order": self.options.get("cache_dit_taylorseer_order", 1),
                    "scm_steps_mask_policy": self.options.get("cache_dit_scm_steps_mask_policy"),
                    "scm_steps_policy": self.options.get("cache_dit_scm_steps_policy", "dynamic"),
                }
            elif cache_backend == "tea_cache":
                cache_config = {
                    "rel_l1_thresh": self.options.get("tea_cache_rel_l1_thresh", 0.2),
                }

            # Base Omni initialization parameters
            omni_kwargs = {
                "model": request.Model,
            }

            # Add diffusion-specific parameters (image/video models)
            if self.model_type in ["image", "video"]:
                omni_kwargs.update({
                    "vae_use_slicing": is_npu(),
                    "vae_use_tiling": is_npu(),
                    "cache_backend": cache_backend,
                    "cache_config": cache_config,
                    "parallel_config": parallel_config,
                    "enforce_eager": self.options.get("enforce_eager", request.EnforceEager),
                    "enable_cpu_offload": self.options.get("enable_cpu_offload", False),
                })
                # Video-specific parameters
                if self.model_type == "video":
                    omni_kwargs.update({
                        "boundary_ratio": self.options.get("boundary_ratio", 0.875),
                        "flow_shift": self.options.get("flow_shift", 5.0),
                    })

            # Add LLM/TTS-specific parameters
            if self.model_type in ["llm", "tts"]:
                omni_kwargs.update({
                    "stage_configs_path": self.options.get("stage_configs_path"),
                    "log_stats": self.options.get("enable_stats", False),
                    "stage_init_timeout": self.options.get("stage_init_timeout", 300),
                })
                # vllm engine options (passed through Omni for LLM/TTS)
                if request.GPUMemoryUtilization > 0:
                    omni_kwargs["gpu_memory_utilization"] = request.GPUMemoryUtilization
                if request.TensorParallelSize > 0:
                    omni_kwargs["tensor_parallel_size"] = request.TensorParallelSize
                if request.TrustRemoteCode:
                    omni_kwargs["trust_remote_code"] = request.TrustRemoteCode
                if request.MaxModelLen > 0:
                    omni_kwargs["max_model_len"] = request.MaxModelLen

            self.omni = Omni(**omni_kwargs)

            # Load tokenizer for LLM/TTS so chat templates work
            if self.model_type in ("llm", "tts"):
                try:
                    from vllm.transformers_utils.tokenizer import get_tokenizer
                    self.tokenizer = get_tokenizer(
                        request.Model,
                        trust_remote_code=opts.get("trust_remote_code", False),
                    )
                except Exception as e:
                    print(f"Failed to load tokenizer: {e}", file=sys.stderr)
                    self.tokenizer = None
            else:
                self.tokenizer = None

            # Setup optional tool / reasoning parsers
            self.tool_parser_cls, self.reasoning_parser_cls = setup_parsers(opts)

            print("Model loaded successfully", file=sys.stderr)
            return backend_pb2.Result(message="Model loaded successfully", success=True)

        except Exception as err:
            print(f"Unexpected {err=}, {type(err)=}", file=sys.stderr)
            traceback.print_exc()
            return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")

    def GenerateImage(self, request, context):
        try:
            # Validate model is loaded and is image/diffusion type
            if not hasattr(self, 'omni'):
                return backend_pb2.Result(success=False, message="Model not loaded. Call LoadModel first.")
            if self.model_type not in ["image"]:
                return backend_pb2.Result(success=False, message=f"Model type {self.model_type} does not support image generation")

            # Extract parameters
            prompt = request.positive_prompt
            negative_prompt = request.negative_prompt if request.negative_prompt else None
            width = request.width if request.width > 0 else 1024
            height = request.height if request.height > 0 else 1024
            seed = request.seed if request.seed > 0 else None
            num_inference_steps = request.step if request.step > 0 else 50
            cfg_scale = self.options.get("cfg_scale", 4.0)
            guidance_scale = self.options.get("guidance_scale", 1.0)

            # Create generator if seed provided
            generator = None
            if seed:
                device = detect_device_type()
                generator = torch.Generator(device=device).manual_seed(seed)

            # Handle image input for image editing
            pil_image = None
            if request.src or (request.ref_images and len(request.ref_images) > 0):
                image_path = request.ref_images[0] if request.ref_images else request.src
                pil_image = self._load_image(image_path)
                if pil_image is None:
                    return backend_pb2.Result(success=False, message=f"Invalid image source: {image_path}")
                pil_image = pil_image.convert("RGB")

            # Build generate kwargs
            generate_kwargs = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "height": height,
                "width": width,
                "generator": generator,
                "true_cfg_scale": cfg_scale,
                "guidance_scale": guidance_scale,
                "num_inference_steps": num_inference_steps,
            }
            if pil_image:
                generate_kwargs["pil_image"] = pil_image

            # Call omni.generate()
            outputs = self.omni.generate(**generate_kwargs)

            # Extract images (following example pattern)
            if not outputs or len(outputs) == 0:
                return backend_pb2.Result(success=False, message="No output generated")

            first_output = outputs[0]
            if not hasattr(first_output, "request_output") or not first_output.request_output:
                return backend_pb2.Result(success=False, message="Invalid output structure")

            req_out = first_output.request_output[0]
            if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
                return backend_pb2.Result(success=False, message="No images in output")

            images = req_out.images
            if not images or len(images) == 0:
                return backend_pb2.Result(success=False, message="Empty images list")

            # Save image
            output_image = images[0]
            output_image.save(request.dst)
            return backend_pb2.Result(message="Image generated successfully", success=True)

        except Exception as err:
            print(f"Error generating image: {err}", file=sys.stderr)
            traceback.print_exc()
            return backend_pb2.Result(success=False, message=f"Error generating image: {err}")

    def GenerateVideo(self, request, context):
        try:
            # Validate model is loaded and is video/diffusion type
            if not hasattr(self, 'omni'):
                return backend_pb2.Result(success=False, message="Model not loaded. Call LoadModel first.")
            if self.model_type not in ["video"]:
                return backend_pb2.Result(success=False, message=f"Model type {self.model_type} does not support video generation")

            # Extract parameters
            prompt = request.prompt
            negative_prompt = request.negative_prompt if request.negative_prompt else ""
            width = request.width if request.width > 0 else 1280
            height = request.height if request.height > 0 else 720
            num_frames = request.num_frames if request.num_frames > 0 else 81
            fps = request.fps if request.fps > 0 else 24
            seed = request.seed if request.seed > 0 else None
            guidance_scale = request.cfg_scale if request.cfg_scale > 0 else 4.0
            guidance_scale_high = self.options.get("guidance_scale_high")
            num_inference_steps = request.step if request.step > 0 else 40

            # Create generator
            generator = None
            if seed:
                device = detect_device_type()
                generator = torch.Generator(device=device).manual_seed(seed)

            # Handle image input for image-to-video
            pil_image = None
            if request.start_image:
                pil_image = self._load_image(request.start_image)
                if pil_image is None:
                    return backend_pb2.Result(success=False, message=f"Invalid start_image: {request.start_image}")
                pil_image = pil_image.convert("RGB")
                # Resize to target dimensions
                pil_image = pil_image.resize((width, height), Image.Resampling.LANCZOS)

            # Build generate kwargs
            generate_kwargs = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "height": height,
                "width": width,
                "generator": generator,
                "guidance_scale": guidance_scale,
                "num_inference_steps": num_inference_steps,
                "num_frames": num_frames,
            }
            if pil_image:
                generate_kwargs["pil_image"] = pil_image
            if guidance_scale_high:
                generate_kwargs["guidance_scale_2"] = guidance_scale_high

            # Call omni.generate()
            frames = self.omni.generate(**generate_kwargs)

            # Extract video frames (following example pattern)
            if isinstance(frames, list) and len(frames) > 0:
                first_item = frames[0]

                if hasattr(first_item, "final_output_type"):
                    if first_item.final_output_type != "image":
                        return backend_pb2.Result(success=False, message=f"Unexpected output type: {first_item.final_output_type}")

                    # Pipeline mode: extract from nested request_output
                    if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
                        if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
                            inner_output = first_item.request_output[0]
                            if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
                                frames = inner_output.images[0] if inner_output.images else None
                    # Diffusion mode: use direct images field
                    elif hasattr(first_item, "images") and first_item.images:
                        frames = first_item.images
                    else:
                        return backend_pb2.Result(success=False, message="No video frames found")

            if frames is None:
                return backend_pb2.Result(success=False, message="No video frames found in output")

            # Convert frames to numpy array (following example)
            if isinstance(frames, torch.Tensor):
                video_tensor = frames.detach().cpu()
                # Handle different tensor shapes [B, C, F, H, W] or [B, F, H, W, C]
                if video_tensor.dim() == 5:
                    if video_tensor.shape[1] in (3, 4):
                        video_tensor = video_tensor[0].permute(1, 2, 3, 0)
                    else:
                        video_tensor = video_tensor[0]
                elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
                    video_tensor = video_tensor.permute(1, 2, 3, 0)
                # Normalize from [-1,1] to [0,1] if float
                if video_tensor.is_floating_point():
                    video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
                video_array = video_tensor.float().numpy()
            else:
                video_array = frames
                if hasattr(video_array, "shape") and video_array.ndim == 5:
                    video_array = video_array[0]

            # Convert 4D array (frames, H, W, C) to list of frames
            if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
                video_array = list(video_array)

            # Save video
            export_to_video(video_array, request.dst, fps=fps)
            return backend_pb2.Result(message="Video generated successfully", success=True)

        except Exception as err:
            print(f"Error generating video: {err}", file=sys.stderr)
            traceback.print_exc()
            return backend_pb2.Result(success=False, message=f"Error generating video: {err}")

    def Predict(self, request, context):
        """Non-streaming text generation with multimodal inputs."""
        gen = self._predict(request, context, streaming=False)
        try:
            res = next(gen)
            return res
        except StopIteration:
            return backend_pb2.Reply(message=bytes("", 'utf-8'))

    def PredictStream(self, request, context):
        """Streaming text generation with multimodal inputs."""
        return self._predict(request, context, streaming=True)

    def _predict(self, request, context, streaming=False):
        """Internal method for text generation (streaming and non-streaming)."""
        try:
            # Validate model is loaded and is LLM type
            if not hasattr(self, 'omni'):
                yield backend_pb2.Reply(message=bytes("Model not loaded. Call LoadModel first.", 'utf-8'))
                return
            if self.model_type not in ["llm"]:
                yield backend_pb2.Reply(message=bytes(f"Model type {self.model_type} does not support text generation", 'utf-8'))
                return

            # Extract prompt
            if request.Prompt:
                prompt = request.Prompt
            elif request.Messages:
                if getattr(self, "tokenizer", None) is not None:
                    messages_dicts = messages_to_dicts(request.Messages)
                    template_kwargs = {"tokenize": False, "add_generation_prompt": True}
                    if request.Tools:
                        try:
                            template_kwargs["tools"] = json.loads(request.Tools)
                        except json.JSONDecodeError:
                            pass
                    try:
                        if request.Metadata.get("enable_thinking", "").lower() == "true":
                            template_kwargs["enable_thinking"] = True
                    except Exception:
                        pass
                    try:
                        prompt = self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
                    except TypeError:
                        prompt = self.tokenizer.apply_chat_template(
                            messages_dicts, tokenize=False, add_generation_prompt=True
                        )
                else:
                    # Fallback: basic template
                    prompt = ""
                    for msg in request.Messages:
                        prompt += f"<|im_start|>{msg.role}\n{msg.content}<|im_end|>\n"
                    prompt += "<|im_start|>assistant\n"
            else:
                yield backend_pb2.Reply(message=bytes("", 'utf-8'))
                return

            # Build multi_modal_data dict
            multi_modal_data = {}

            # Process images
            if request.Images:
                image_data = []
                for img_path in request.Images:
                    img = self._load_image(img_path)
                    if img:
                        # Convert to format expected by vllm
                        from vllm.multimodal.image import convert_image_mode
                        img_data = convert_image_mode(img, "RGB")
                        image_data.append(img_data)
                if image_data:
                    multi_modal_data["image"] = image_data

            # Process videos
            if request.Videos:
                video_data = []
                for video_path in request.Videos:
                    video = self._load_video(video_path)
                    if video is not None:
                        video_data.append(video)
                if video_data:
                    multi_modal_data["video"] = video_data

            # Process audio
            if request.Audios:
                audio_data = []
                for audio_path in request.Audios:
                    audio = self._load_audio(audio_path)
                    if audio is not None:
                        audio_data.append(audio)
                if audio_data:
                    multi_modal_data["audio"] = audio_data

            # Build inputs dict
            inputs = {
                "prompt": prompt,
                "multi_modal_data": multi_modal_data if multi_modal_data else None,
            }

            # Build sampling params
            sampling_params = SamplingParams(
                temperature=request.Temperature if request.Temperature > 0 else 0.7,
                top_p=request.TopP if request.TopP > 0 else 0.9,
                top_k=request.TopK if request.TopK > 0 else -1,
                max_tokens=request.Tokens if request.Tokens > 0 else 200,
                presence_penalty=request.PresencePenalty if request.PresencePenalty != 0 else 0.0,
                frequency_penalty=request.FrequencyPenalty if request.FrequencyPenalty != 0 else 0.0,
                repetition_penalty=request.RepetitionPenalty if request.RepetitionPenalty != 0 else 1.0,
                seed=request.Seed if request.Seed > 0 else None,
                stop=request.StopPrompts if request.StopPrompts else None,
                stop_token_ids=request.StopTokenIds if request.StopTokenIds else None,
                ignore_eos=request.IgnoreEOS,
            )
            sampling_params_list = [sampling_params]

            # Call omni.generate() (returns generator for LLM mode)
            omni_generator = self.omni.generate([inputs], sampling_params_list)

            # Extract text from outputs and track token usage
            generated_text = ""
            prompt_tokens = 0
            completion_tokens = 0
            for stage_outputs in omni_generator:
                if stage_outputs.final_output_type == "text":
                    for output in stage_outputs.request_output:
                        completion = output.outputs[0]
                        text_output = completion.text
                        # Track tokens when available
                        try:
                            if getattr(output, "prompt_token_ids", None) is not None:
                                prompt_tokens = len(output.prompt_token_ids)
                            if getattr(completion, "token_ids", None) is not None:
                                completion_tokens = len(completion.token_ids)
                        except Exception:
                            pass
                        if streaming:
                            # Remove already sent text (vllm concatenates)
                            delta_text = text_output.removeprefix(generated_text)
                            yield backend_pb2.Reply(
                                message=bytes(delta_text, encoding='utf-8'),
                                tokens=completion_tokens,
                                prompt_tokens=prompt_tokens,
                            )
                        generated_text = text_output

            if not streaming:
                # Build optional ChatDelta with parsed reasoning / tool calls
                chat_deltas = []
                content_text = generated_text
                reasoning_text = ""
                tool_call_deltas = []

                if self.reasoning_parser_cls is not None:
                    try:
                        parser = self.reasoning_parser_cls(self.tokenizer) if self.tokenizer else self.reasoning_parser_cls()
                        reasoning_text, content_text = parser.extract_reasoning_content(content_text, request=None)
                        reasoning_text = reasoning_text or ""
                        content_text = content_text or ""
                    except Exception as e:
                        print(f"reasoning_parser failed: {e}", file=sys.stderr)

                if self.tool_parser_cls is not None:
                    try:
                        parser = self.tool_parser_cls(self.tokenizer) if self.tokenizer else self.tool_parser_cls()
                        tool_info = parser.extract_tool_calls(content_text, request=None)
                        if getattr(tool_info, "tools_called", False):
                            content_text = tool_info.content or ""
                            for tc in tool_info.tool_calls or []:
                                fn = getattr(tc, "function", None)
                                tool_call_deltas.append(backend_pb2.ToolCallDelta(
                                    index=getattr(tc, "index", 0) or 0,
                                    id=getattr(tc, "id", "") or "",
                                    name=getattr(fn, "name", "") if fn else "",
                                    arguments=getattr(fn, "arguments", "") if fn else "",
                                ))
                    except Exception as e:
                        print(f"tool_parser failed: {e}", file=sys.stderr)

                if self.tool_parser_cls is not None or self.reasoning_parser_cls is not None:
                    chat_deltas.append(backend_pb2.ChatDelta(
                        content=content_text,
                        reasoning_content=reasoning_text,
                        tool_calls=tool_call_deltas,
                    ))

                yield backend_pb2.Reply(
                    message=bytes(generated_text, encoding='utf-8'),
                    tokens=completion_tokens,
                    prompt_tokens=prompt_tokens,
                    chat_deltas=chat_deltas,
                )

        except Exception as err:
            print(f"Error in Predict: {err}", file=sys.stderr)
            traceback.print_exc()
            yield backend_pb2.Reply(message=bytes(f"Error: {err}", encoding='utf-8'))

    def TTS(self, request, context):
        try:
            # Validate model is loaded and is TTS type
            if not hasattr(self, 'omni'):
                return backend_pb2.Result(success=False, message="Model not loaded. Call LoadModel first.")
            if self.model_type not in ["tts"]:
                return backend_pb2.Result(success=False, message=f"Model type {self.model_type} does not support TTS")

            # Extract parameters
            text = request.text
            language = request.language if request.language else "Auto"
            voice = request.voice if request.voice else None
            task_type = self._detect_tts_task_type()

            # Build prompt with chat template
            # TODO: for now vllm-omni supports only qwen3-tts, so we hardcode it, however, we want to support other models in the future.
            # and we might need to use the chat template here
            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"

            # Build inputs dict
            inputs = {
                "prompt": prompt,
                "additional_information": {
                    "task_type": [task_type],
                    "text": [text],
                    "language": [language],
                    "max_new_tokens": [2048],
                }
            }

            # Add task-specific fields
            if task_type == "CustomVoice":
                if voice:
                    inputs["additional_information"]["speaker"] = [voice]
                # Add instruct if provided in options
                if "instruct" in self.options:
                    inputs["additional_information"]["instruct"] = [self.options["instruct"]]
            elif task_type == "VoiceDesign":
                if "instruct" in self.options:
                    inputs["additional_information"]["instruct"] = [self.options["instruct"]]
                inputs["additional_information"]["non_streaming_mode"] = [True]
            elif task_type == "Base":
                # Voice cloning requires ref_audio and ref_text
                if "ref_audio" in self.options:
                    inputs["additional_information"]["ref_audio"] = [self.options["ref_audio"]]
                if "ref_text" in self.options:
                    inputs["additional_information"]["ref_text"] = [self.options["ref_text"]]
                if "x_vector_only_mode" in self.options:
                    inputs["additional_information"]["x_vector_only_mode"] = [self.options["x_vector_only_mode"]]

            # Build sampling params
            sampling_params = SamplingParams(
                temperature=0.9,
                top_p=1.0,
                top_k=50,
                max_tokens=2048,
                seed=42,
                detokenize=False,
                repetition_penalty=1.05,
            )
            sampling_params_list = [sampling_params]

            # Call omni.generate()
            omni_generator = self.omni.generate(inputs, sampling_params_list)

            # Extract audio (following TTS example)
            for stage_outputs in omni_generator:
                for output in stage_outputs.request_output:
                    if "audio" in output.multimodal_output:
                        audio_tensor = output.multimodal_output["audio"]
                        audio_samplerate = output.multimodal_output["sr"].item()

                        # Convert to numpy
                        audio_numpy = audio_tensor.float().detach().cpu().numpy()
                        if audio_numpy.ndim > 1:
                            audio_numpy = audio_numpy.flatten()

                        # Save audio file
                        sf.write(request.dst, audio_numpy, samplerate=audio_samplerate, format="WAV")
                        return backend_pb2.Result(message="TTS audio generated successfully", success=True)

            return backend_pb2.Result(success=False, message="No audio output generated")

        except Exception as err:
            print(f"Error generating TTS: {err}", file=sys.stderr)
            traceback.print_exc()
            return backend_pb2.Result(success=False, message=f"Error generating TTS: {err}")

    def TokenizeString(self, request, context):
        if not hasattr(self, 'tokenizer') or self.tokenizer is None:
            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
            context.set_details("Model/tokenizer not loaded")
            return backend_pb2.TokenizationResponse()
        try:
            tokens = self.tokenizer.encode(request.Prompt)
            return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
        except Exception as e:
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(str(e))
            return backend_pb2.TokenizationResponse()

    def Free(self, request, context):
        try:
            if hasattr(self, 'omni'):
                del self.omni
            if hasattr(self, 'tokenizer'):
                del self.tokenizer
            self.tool_parser_cls = None
            self.reasoning_parser_cls = None
            gc.collect()
            try:
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except Exception:
                pass
            return backend_pb2.Result(success=True, message="Model freed")
        except Exception as e:
            return backend_pb2.Result(success=False, message=str(e))


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),
            ('grpc.max_receive_message_length', 50 * 1024 * 1024),
        ],
        interceptors=get_auth_interceptors(),
    )
    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)

    # Signal handlers for graceful shutdown
    def signal_handler(sig, frame):
        print("Received termination signal. Shutting down...")
        server.stop(0)
        sys.exit(0)

    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)