improve

src/server/main-v2.py

@@ -5,7 +5,7 @@ from time import time
 from typing import List, Dict
 import tempfile
 import uvicorn
-from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Body, Form, APIRouter
+from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Body, Form
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse
 from PIL import Image
@@ -22,17 +22,18 @@ from contextlib import asynccontextmanager
 import soundfile as sf
 import numpy as np
 import requests
+import logging
 from starlette.responses import StreamingResponse
-from logging_config import logger
-from tts_config import SPEED, ResponseFormat, config as tts_config
+from logging_config import logger  # Assumed external logging config
+from tts_config import SPEED, ResponseFormat, config as tts_config  # Assumed external TTS config
 import torchaudio
 from tenacity import retry, stop_after_attempt, wait_exponential
 from torch.cuda.amp import autocast
 
 # Device setup
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.bfloat16 if device != "cpu" else torch.float32
-logger.info(f"{'GPU' if device != 'cpu' else 'CPU'} will be used for inference")
+torch_dtype = torch.float16 if device != "cpu" else torch.float32
+logger.info(f"Using device: {device} with dtype: {torch_dtype}")
 
 # Check CUDA availability and version
 cuda_available = torch.cuda.is_available()
@@ -77,12 +78,15 @@ quantization_config = BitsAndBytesConfig(
 # Request queue for concurrency control
 request_queue = asyncio.Queue(maxsize=10)
 
+# Logging optimization
+logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
+
 # LLM Manager with batching
 class LLMManager:
     def __init__(self, model_name: str, device: str = device):
         self.model_name = model_name
         self.device = torch.device(device)
-        self.torch_dtype = torch.bfloat16 if self.device.type != "cpu" else torch.float32
+        self.torch_dtype = torch.float16 if self.device.type != "cpu" else torch.float32
         self.model = None
         self.processor = None
         self.is_loaded = False
@@ -93,6 +97,9 @@ class LLMManager:
     def load(self):
         if not self.is_loaded:
             try:
+                if self.device.type == "cuda":
+                    torch.set_float32_matmul_precision('high')
+                    logger.info("Enabled TF32 matrix multiplication for improved GPU performance")
                 self.model = Gemma3ForConditionalGeneration.from_pretrained(
                     self.model_name,
                     device_map="auto",
@@ -107,7 +114,7 @@ class LLMManager:
                 logger.info(f"LLM {self.model_name} loaded and warmed up on {self.device}")
             except Exception as e:
                 logger.error(f"Failed to load LLM: {str(e)}")
-                raise
+                self.is_loaded = False
 
     def unload(self):
         if self.is_loaded:
@@ -115,14 +122,17 @@ class LLMManager:
             del self.processor
             if self.device.type == "cuda":
                 torch.cuda.empty_cache()
-                logger.info(f"GPU memory allocated after unload: {torch.cuda.memory_allocated()}")
+                logger.info(f"GPU memory cleared: {torch.cuda.memory_allocated()} bytes allocated")
             self.is_loaded = False
             self.token_cache.clear()
-            logger.info(f"LLM {self.model_name} unloaded from {self.device}")
+            logger.info(f"LLM {self.model_name} unloaded")
 
     async def generate(self, prompt: str, max_tokens: int = settings.max_tokens, temperature: float = 0.7) -> str:
         if not self.is_loaded:
+            logger.warning("LLM not loaded; attempting reload")
             self.load()
+        if not self.is_loaded:
+            raise HTTPException(status_code=503, detail="LLM model unavailable")
 
         cache_key = f"{prompt}:{max_tokens}:{temperature}"
         if cache_key in self.token_cache:
@@ -153,7 +163,6 @@ class LLMManager:
                 return_tensors="pt",
                 padding=True
             ).to(self.device, dtype=torch.bfloat16)
-
             with autocast(), torch.no_grad():
                 outputs = self.model.generate(
                     **inputs_vlm,
@@ -175,12 +184,10 @@ class LLMManager:
     async def vision_query(self, image: Image.Image, query: str) -> str:
         if not self.is_loaded:
             self.load()
-
         messages_vlm = [
             {"role": "system", "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Summarize your answer in maximum 1 sentence."}]},
             {"role": "user", "content": [{"type": "text", "text": query}] + ([{"type": "image", "image": image}] if image and image.size[0] > 0 and image.size[1] > 0 else [])}
         ]
-
         try:
             inputs_vlm = self.processor.apply_chat_template(
                 messages_vlm,
@@ -192,7 +199,6 @@ class LLMManager:
         except Exception as e:
             logger.error(f"Error in apply_chat_template: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Failed to process input: {str(e)}")
-
         input_len = inputs_vlm["input_ids"].shape[-1]
         with torch.inference_mode():
             generation = self.model.generate(**inputs_vlm, max_new_tokens=512, do_sample=True, temperature=0.7)
@@ -204,12 +210,10 @@ class LLMManager:
     async def chat_v2(self, image: Image.Image, query: str) -> str:
         if not self.is_loaded:
             self.load()
-
         messages_vlm = [
             {"role": "system", "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Answer questions considering India as base country and Karnataka as base state."}]},
             {"role": "user", "content": [{"type": "text", "text": query}] + ([{"type": "image", "image": image}] if image and image.size[0] > 0 and image.size[1] > 0 else [])}
         ]
-
         try:
             inputs_vlm = self.processor.apply_chat_template(
                 messages_vlm,
@@ -221,7 +225,6 @@ class LLMManager:
         except Exception as e:
             logger.error(f"Error in apply_chat_template: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Failed to process input: {str(e)}")
-
         input_len = inputs_vlm["input_ids"].shape[-1]
         with torch.inference_mode():
             generation = self.model.generate(**inputs_vlm, max_new_tokens=512, do_sample=True, temperature=0.7)
@@ -240,9 +243,18 @@ class TTSManager:
 
     def load(self):
         if not self.model:
-            logger.info("Loading TTS model IndicF5...")
+            logger.info(f"Loading TTS model {self.repo_id} on {self.device_type}...")
             self.model = AutoModel.from_pretrained(self.repo_id, trust_remote_code=True).to(self.device_type)
-            logger.info("TTS model IndicF5 loaded")
+            logger.info("TTS model loaded")
+
+    def unload(self):
+        if self.model:
+            del self.model
+            if self.device_type.type == "cuda":
+                torch.cuda.empty_cache()
+                logger.info(f"TTS GPU memory cleared: {torch.cuda.memory_allocated()} bytes allocated")
+            self.model = None
+            logger.info("TTS model unloaded")
 
     def synthesize(self, text, ref_audio_path, ref_text):
         if not self.model:
@@ -320,13 +332,22 @@ class ASRModelManager:
 
     def load(self):
         if not self.model:
-            logger.info("Loading ASR model...")
+            logger.info(f"Loading ASR model on {self.device_type}...")
             self.model = AutoModel.from_pretrained(
                 "ai4bharat/indic-conformer-600m-multilingual",
                 trust_remote_code=True
             ).to(self.device_type)
             logger.info("ASR model loaded")
 
+    def unload(self):
+        if self.model:
+            del self.model
+            if self.device_type.type == "cuda":
+                torch.cuda.empty_cache()
+                logger.info(f"ASR GPU memory cleared: {torch.cuda.memory_allocated()} bytes allocated")
+            self.model = None
+            logger.info("ASR model unloaded")
+
 # Global Managers
 llm_manager = LLMManager(settings.llm_model_name)
 model_manager = ModelManager()
@@ -353,6 +374,12 @@ class SynthesizeRequest(BaseModel):
 class KannadaSynthesizeRequest(BaseModel):
     text: str
 
+    @field_validator("text")
+    def text_must_be_valid(cls, v):
+        if len(v) > 500:
+            raise ValueError("Text cannot exceed 500 characters")
+        return v.strip()
+
 class ChatRequest(BaseModel):
     prompt: str
     src_lang: str = "kan_Knda"
@@ -441,9 +468,6 @@ async def perform_internal_translation(sentences: List[str], src_lang: str, tgt_
         model_manager.load_model(src_lang, tgt_lang, key)
         translate_manager = model_manager.get_model(src_lang, tgt_lang)
 
-    if not translate_manager.model:
-        translate_manager.load()
-
     batch = ip.preprocess_batch(sentences, src_lang=src_lang, tgt_lang=tgt_lang)
     inputs = translate_manager.tokenizer(batch, truncation=True, padding="longest", return_tensors="pt", return_attention_mask=True).to(translate_manager.device_type)
     with torch.no_grad(), autocast():
@@ -479,13 +503,17 @@ async def lifespan(app: FastAPI):
             model_manager.load_model(src_lang, tgt_lang, key)
         logger.info("All models loaded successfully")
 
-    logger.info("Starting sequential model loading...")
+    logger.info("Starting server with preloaded models...")
     load_all_models()
     batch_task = asyncio.create_task(batch_worker())
     yield
     batch_task.cancel()
     llm_manager.unload()
-    logger.info("Server shutdown complete")
+    tts_manager.unload()
+    asr_manager.unload()
+    for model in model_manager.models.values():
+        model.unload()
+    logger.info("Server shutdown complete; all models unloaded")
 
 # Batch Worker
 async def batch_worker():
@@ -519,8 +547,8 @@ async def batch_worker():
 
 # FastAPI App
 app = FastAPI(
-    title="Dhwani API",
-    description="AI Chat API supporting Indian languages",
+    title="Optimized Dhwani API",
+    description="AI Chat API supporting Indian languages with performance enhancements",
     version="1.0.0",
     redirect_slashes=False,
     lifespan=lifespan
@@ -577,6 +605,8 @@ async def translate(request: TranslationRequest, translate_manager: TranslateMan
 @app.get("/v1/health")
 async def health_check():
     memory_usage = torch.cuda.memory_allocated() / (24 * 1024**3) if cuda_available else 0
+    if memory_usage > 0.9:
+        logger.warning("GPU memory usage exceeds 90%; consider unloading models")
     llm_status = "unhealthy"
     llm_latency = None
     if llm_manager.is_loaded:
@@ -587,7 +617,6 @@ async def health_check():
             llm_status = "healthy" if llm_test else "unhealthy"
         except Exception as e:
             logger.error(f"LLM health check failed: {str(e)}")
-
     tts_status = "unhealthy"
     tts_latency = None
     if tts_manager.model:
@@ -598,7 +627,6 @@ async def health_check():
             tts_status = "healthy" if audio_buffer else "unhealthy"
         except Exception as e:
             logger.error(f"TTS health check failed: {str(e)}")
-
     asr_status = "unhealthy"
     asr_latency = None
     if asr_manager.model:
@@ -612,7 +640,6 @@ async def health_check():
             asr_status = "healthy" if asr_test else "unhealthy"
         except Exception as e:
             logger.error(f"ASR health check failed: {str(e)}")
-
     status = {
         "status": "healthy" if llm_status == "healthy" and tts_status == "healthy" and asr_status == "healthy" else "degraded",
         "model": settings.llm_model_name,
@@ -622,6 +649,7 @@ async def health_check():
         "tts_latency": f"{tts_latency:.3f}s" if tts_latency else "N/A",
         "asr_status": asr_status,
         "asr_latency": f"{asr_latency:.3f}s" if asr_latency else "N/A",
+        "translation_models": list(model_manager.models.keys()),
         "gpu_memory_usage": f"{memory_usage:.2%}"
     }
     logger.info("Health check completed")
@@ -636,6 +664,10 @@ async def unload_all_models():
     try:
         logger.info("Starting to unload all models...")
         llm_manager.unload()
+        tts_manager.unload()
+        asr_manager.unload()
+        for model in model_manager.models.values():
+            model.unload()
         logger.info("All models unloaded successfully")
         return {"status": "success", "message": "All models unloaded"}
     except Exception as e:
@@ -647,6 +679,15 @@ async def load_all_models():
     try:
         logger.info("Starting to load all models...")
         llm_manager.load()
+        tts_manager.load()
+        asr_manager.load()
+        for src_lang, tgt_lang, key in [
+            ('eng_Latn', 'kan_Knda', 'eng_indic'),
+            ('kan_Knda', 'eng_Latn', 'indic_eng'),
+            ('kan_Knda', 'hin_Deva', 'indic_indic'),
+        ]:
+            if key not in model_manager.models:
+                model_manager.load_model(src_lang, tgt_lang, key)
         logger.info("All models loaded successfully")
         return {"status": "success", "message": "All models loaded"}
     except Exception as e:
@@ -667,33 +708,32 @@ async def translate_endpoint(request: TranslationRequest):
 @app.post("/v1/chat", response_model=ChatResponse)
 @limiter.limit(settings.chat_rate_limit)
 async def chat(request: Request, chat_request: ChatRequest):
-    if not chat_request.prompt:
-        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-    logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, tgt_lang: {chat_request.tgt_lang}")
-    EUROPEAN_LANGUAGES = {"deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn", "por_Latn", "rus_Cyrl", "pol_Latn"}
-    try:
-        if chat_request.src_lang != "eng_Latn" and chat_request.src_lang not in EUROPEAN_LANGUAGES:
-            translated_prompt = await perform_internal_translation([chat_request.prompt], chat_request.src_lang, "eng_Latn")
-            prompt_to_process = translated_prompt[0]
-            logger.info(f"Translated prompt to English: {prompt_to_process}")
-        else:
-            prompt_to_process = chat_request.prompt
-            logger.info("Prompt in English or European language, no translation needed")
-
-        response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-        logger.info(f"Generated English response: {response}")
-
-        if chat_request.tgt_lang != "eng_Latn" and chat_request.tgt_lang not in EUROPEAN_LANGUAGES:
-            translated_response = await perform_internal_translation([response], "eng_Latn", chat_request.tgt_lang)
-            final_response = translated_response[0]
-            logger.info(f"Translated response to {chat_request.tgt_lang}: {final_response}")
-        else:
-            final_response = response
-            logger.info(f"Response in {chat_request.tgt_lang}, no translation needed")
-        return ChatResponse(response=final_response)
-    except Exception as e:
-        logger.error(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    async with request_queue:
+        if not chat_request.prompt:
+            raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, tgt_lang: {chat_request.tgt_lang}")
+        EUROPEAN_LANGUAGES = {"deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn", "por_Latn", "rus_Cyrl", "pol_Latn"}
+        try:
+            if chat_request.src_lang != "eng_Latn" and chat_request.src_lang not in EUROPEAN_LANGUAGES:
+                translated_prompt = await perform_internal_translation([chat_request.prompt], chat_request.src_lang, "eng_Latn")
+                prompt_to_process = translated_prompt[0]
+                logger.info(f"Translated prompt to English: {prompt_to_process}")
+            else:
+                prompt_to_process = chat_request.prompt
+                logger.info("Prompt in English or European language, no translation needed")
+            response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
+            logger.info(f"Generated English response: {response}")
+            if chat_request.tgt_lang != "eng_Latn" and chat_request.tgt_lang not in EUROPEAN_LANGUAGES:
+                translated_response = await perform_internal_translation([response], "eng_Latn", chat_request.tgt_lang)
+                final_response = translated_response[0]
+                logger.info(f"Translated response to {chat_request.tgt_lang}: {final_response}")
+            else:
+                final_response = response
+                logger.info(f"Response in {chat_request.tgt_lang}, no translation needed")
+            return ChatResponse(response=final_response)
+        except Exception as e:
+            logger.error(f"Error processing request: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 @app.post("/v1/visual_query/")
 async def visual_query(
@@ -702,30 +742,31 @@ async def visual_query(
     src_lang: str = Query("kan_Knda", enum=list(SUPPORTED_LANGUAGES)),
     tgt_lang: str = Query("kan_Knda", enum=list(SUPPORTED_LANGUAGES)),
 ):
-    try:
-        image = Image.open(file.file)
-        if image.size == (0, 0):
-            raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
-        if src_lang != "eng_Latn":
-            translated_query = await perform_internal_translation([query], src_lang, "eng_Latn")
-            query_to_process = translated_query[0]
-            logger.info(f"Translated query to English: {query_to_process}")
-        else:
-            query_to_process = query
-            logger.info("Query already in English, no translation needed")
-        answer = await llm_manager.vision_query(image, query_to_process)
-        logger.info(f"Generated English answer: {answer}")
-        if tgt_lang != "eng_Latn":
-            translated_answer = await perform_internal_translation([answer], "eng_Latn", tgt_lang)
-            final_answer = translated_answer[0]
-            logger.info(f"Translated answer to {tgt_lang}: {final_answer}")
-        else:
-            final_answer = answer
-            logger.info("Answer kept in English, no translation needed")
-        return {"answer": final_answer}
-    except Exception as e:
-        logger.error(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    async with request_queue:
+        try:
+            image = Image.open(file.file)
+            if image.size == (0, 0):
+                raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
+            if src_lang != "eng_Latn":
+                translated_query = await perform_internal_translation([query], src_lang, "eng_Latn")
+                query_to_process = translated_query[0]
+                logger.info(f"Translated query to English: {query_to_process}")
+            else:
+                query_to_process = query
+                logger.info("Query already in English, no translation needed")
+            answer = await llm_manager.vision_query(image, query_to_process)
+            logger.info(f"Generated English answer: {answer}")
+            if tgt_lang != "eng_Latn":
+                translated_answer = await perform_internal_translation([answer], "eng_Latn", tgt_lang)
+                final_answer = translated_answer[0]
+                logger.info(f"Translated answer to {tgt_lang}: {final_answer}")
+            else:
+                final_answer = answer
+                logger.info("Answer kept in English, no translation needed")
+            return {"answer": final_answer}
+        except Exception as e:
+            logger.error(f"Error processing request: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 @app.post("/v1/chat_v2", response_model=ChatResponse)
 @limiter.limit(settings.chat_rate_limit)
@@ -736,68 +777,70 @@ async def chat_v2(
     src_lang: str = Form("kan_Knda"),
     tgt_lang: str = Form("kan_Knda"),
 ):
-    if not prompt:
-        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-    if src_lang not in SUPPORTED_LANGUAGES or tgt_lang not in SUPPORTED_LANGUAGES:
-        raise HTTPException(status_code=400, detail=f"Unsupported language code. Supported codes: {', '.join(SUPPORTED_LANGUAGES)}")
-    logger.info(f"Received prompt: {prompt}, src_lang: {src_lang}, tgt_lang: {tgt_lang}, Image provided: {image is not None}")
-    try:
-        if image:
-            image_data = await image.read()
-            if not image_data:
-                raise HTTPException(status_code=400, detail="Uploaded image is empty")
-            img = Image.open(io.BytesIO(image_data))
-            if src_lang != "eng_Latn":
-                translated_prompt = await perform_internal_translation([prompt], src_lang, "eng_Latn")
-                prompt_to_process = translated_prompt[0]
-                logger.info(f"Translated prompt to English: {prompt_to_process}")
-            else:
-                prompt_to_process = prompt
-            decoded = await llm_manager.chat_v2(img, prompt_to_process)
-            logger.info(f"Generated English response: {decoded}")
-            if tgt_lang != "eng_Latn":
-                translated_response = await perform_internal_translation([decoded], "eng_Latn", tgt_lang)
-                final_response = translated_response[0]
-                logger.info(f"Translated response to {tgt_lang}: {final_response}")
-            else:
-                final_response = decoded
-        else:
-            if src_lang != "eng_Latn":
-                translated_prompt = await perform_internal_translation([prompt], src_lang, "eng_Latn")
-                prompt_to_process = translated_prompt[0]
-                logger.info(f"Translated prompt to English: {prompt_to_process}")
-            else:
-                prompt_to_process = prompt
-            decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-            logger.info(f"Generated English response: {decoded}")
-            if tgt_lang != "eng_Latn":
-                translated_response = await perform_internal_translation([decoded], "eng_Latn", tgt_lang)
-                final_response = translated_response[0]
-                logger.info(f"Translated response to {tgt_lang}: {final_response}")
+    async with request_queue:
+        if not prompt:
+            raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        if src_lang not in SUPPORTED_LANGUAGES or tgt_lang not in SUPPORTED_LANGUAGES:
+            raise HTTPException(status_code=400, detail=f"Unsupported language code. Supported codes: {', '.join(SUPPORTED_LANGUAGES)}")
+        logger.info(f"Received prompt: {prompt}, src_lang: {src_lang}, tgt_lang: {tgt_lang}, Image provided: {image is not None}")
+        try:
+            if image:
+                image_data = await image.read()
+                if not image_data:
+                    raise HTTPException(status_code=400, detail="Uploaded image is empty")
+                img = Image.open(io.BytesIO(image_data))
+                if src_lang != "eng_Latn":
+                    translated_prompt = await perform_internal_translation([prompt], src_lang, "eng_Latn")
+                    prompt_to_process = translated_prompt[0]
+                    logger.info(f"Translated prompt to English: {prompt_to_process}")
+                else:
+                    prompt_to_process = prompt
+                decoded = await llm_manager.chat_v2(img, prompt_to_process)
+                logger.info(f"Generated English response: {decoded}")
+                if tgt_lang != "eng_Latn":
+                    translated_response = await perform_internal_translation([decoded], "eng_Latn", tgt_lang)
+                    final_response = translated_response[0]
+                    logger.info(f"Translated response to {tgt_lang}: {final_response}")
+                else:
+                    final_response = decoded
             else:
-                final_response = decoded
-        return ChatResponse(response=final_response)
-    except Exception as e:
-        logger.error(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+                if src_lang != "eng_Latn":
+                    translated_prompt = await perform_internal_translation([prompt], src_lang, "eng_Latn")
+                    prompt_to_process = translated_prompt[0]
+                    logger.info(f"Translated prompt to English: {prompt_to_process}")
+                else:
+                    prompt_to_process = prompt
+                decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
+                logger.info(f"Generated English response: {decoded}")
+                if tgt_lang != "eng_Latn":
+                    translated_response = await perform_internal_translation([decoded], "eng_Latn", tgt_lang)
+                    final_response = translated_response[0]
+                    logger.info(f"Translated response to {tgt_lang}: {final_response}")
+                else:
+                    final_response = decoded
+            return ChatResponse(response=final_response)
+        except Exception as e:
+            logger.error(f"Error processing request: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 @app.post("/transcribe/", response_model=TranscriptionResponse)
 async def transcribe_audio(file: UploadFile = File(...), language: str = Query(..., enum=list(asr_manager.model_language.keys()))):
-    if not asr_manager.model:
-        raise HTTPException(status_code=503, detail="ASR model not loaded")
-    try:
-        wav, sr = torchaudio.load(file.file, backend="cuda" if cuda_available else "cpu")
-        wav = torch.mean(wav, dim=0, keepdim=True).to(device)
-        target_sample_rate = 16000
-        if sr != target_sample_rate:
-            resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sample_rate).to(device)
-            wav = resampler(wav)
-        with autocast(), torch.no_grad():
-            transcription_rnnt = asr_manager.model(wav, asr_manager.model_language[language], "rnnt")
-        return TranscriptionResponse(text=transcription_rnnt)
-    except Exception as e:
-        logger.error(f"Error in transcription: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
+    async with request_queue:
+        if not asr_manager.model:
+            raise HTTPException(status_code=503, detail="ASR model not loaded")
+        try:
+            wav, sr = torchaudio.load(file.file, backend="cuda" if cuda_available else "cpu")
+            wav = torch.mean(wav, dim=0, keepdim=True).to(device)
+            target_sample_rate = 16000
+            if sr != target_sample_rate:
+                resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sample_rate).to(device)
+                wav = resampler(wav)
+            with autocast(), torch.no_grad():
+                transcription_rnnt = asr_manager.model(wav, asr_manager.model_language[language], "rnnt")
+            return TranscriptionResponse(text=transcription_rnnt)
+        except Exception as e:
+            logger.error(f"Error in transcription: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
 
 @app.post("/v1/speech_to_speech")
 async def speech_to_speech(
@@ -805,16 +848,17 @@ async def speech_to_speech(
     file: UploadFile = File(...),
     language: str = Query(..., enum=list(asr_manager.model_language.keys())),
 ) -> StreamingResponse:
-    if not tts_manager.model:
-        raise HTTPException(status_code=503, detail="TTS model not loaded")
-    transcription = await transcribe_audio(file, language)
-    logger.info(f"Transcribed text: {transcription.text}")
-    chat_request = ChatRequest(prompt=transcription.text, src_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"), tgt_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"))
-    processed_text = await chat(request, chat_request)
-    logger.info(f"Processed text: {processed_text.response}")
-    voice_request = KannadaSynthesizeRequest(text=processed_text.response)
-    audio_response = await synthesize_kannada(voice_request)
-    return audio_response
+    async with request_queue:
+        if not tts_manager.model:
+            raise HTTPException(status_code=503, detail="TTS model not loaded")
+        transcription = await transcribe_audio(file, language)
+        logger.info(f"Transcribed text: {transcription.text}")
+        chat_request = ChatRequest(prompt=transcription.text, src_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"), tgt_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"))
+        processed_text = await chat(request, chat_request)
+        logger.info(f"Processed text: {processed_text.response}")
+        voice_request = KannadaSynthesizeRequest(text=processed_text.response)
+        audio_response = await synthesize_kannada(voice_request)
+        return audio_response
 
 LANGUAGE_TO_SCRIPT = {"kannada": "kan_Knda"}
 

src/server/main.py

@@ -2,7 +2,7 @@ import argparse
 import io
 import os
 from time import time
-from typing import List
+from typing import List, Dict
 import tempfile
 import uvicorn
 from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile, Body, Form
@@ -22,25 +22,23 @@ from contextlib import asynccontextmanager
 import soundfile as sf
 import numpy as np
 import requests
+import logging
 from starlette.responses import StreamingResponse
-from logging_config import logger
-from tts_config import SPEED, ResponseFormat, config as tts_config
+from logging_config import logger  # Assumed external logging config
+from tts_config import SPEED, ResponseFormat, config as tts_config  # Assumed external TTS config
 import torchaudio
+from tenacity import retry, stop_after_attempt, wait_exponential
+from torch.cuda.amp import autocast
 
 # Device setup
-if torch.cuda.is_available():
-    device = "cuda:0"
-    logger.info("GPU will be used for inference")
-else:
-    device = "cpu"
-    logger.info("CPU will be used for inference")
-torch_dtype = torch.bfloat16 if device != "cpu" else torch.float32
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if device != "cpu" else torch.float32
+logger.info(f"Using device: {device} with dtype: {torch_dtype}")
 
 # Check CUDA availability and version
 cuda_available = torch.cuda.is_available()
 cuda_version = torch.version.cuda if cuda_available else None
-
-if torch.cuda.is_available():
+if cuda_available:
     device_idx = torch.cuda.current_device()
     capability = torch.cuda.get_device_capability(device_idx)
     compute_capability_float = float(f"{capability[0]}.{capability[1]}")
@@ -77,33 +75,46 @@ quantization_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-# LLM Manager
+# Request queue for concurrency control
+request_queue = asyncio.Queue(maxsize=10)
+
+# Logging optimization
+logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
+
+# LLM Manager with batching
 class LLMManager:
-    def __init__(self, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
+    def __init__(self, model_name: str, device: str = device):
         self.model_name = model_name
         self.device = torch.device(device)
-        self.torch_dtype = torch.bfloat16 if self.device.type != "cpu" else torch.float32
+        self.torch_dtype = torch.float16 if self.device.type != "cpu" else torch.float32
         self.model = None
         self.processor = None
         self.is_loaded = False
+        self.token_cache = {}
+        self.load()
         logger.info(f"LLMManager initialized with model {model_name} on {self.device}")
 
     def load(self):
         if not self.is_loaded:
             try:
+                if self.device.type == "cuda":
+                    torch.set_float32_matmul_precision('high')
+                    logger.info("Enabled TF32 matrix multiplication for improved GPU performance")
                 self.model = Gemma3ForConditionalGeneration.from_pretrained(
                     self.model_name,
                     device_map="auto",
                     quantization_config=quantization_config,
                     torch_dtype=self.torch_dtype
-                )
-                self.model.eval()
+                ).eval()
                 self.processor = AutoProcessor.from_pretrained(self.model_name)
+                dummy_input = self.processor("test", return_tensors="pt").to(self.device)
+                with torch.no_grad():
+                    self.model.generate(**dummy_input, max_new_tokens=10)
                 self.is_loaded = True
-                logger.info(f"LLM {self.model_name} loaded on {self.device}")
+                logger.info(f"LLM {self.model_name} loaded and warmed up on {self.device}")
             except Exception as e:
                 logger.error(f"Failed to load LLM: {str(e)}")
-                raise
+                self.is_loaded = False
 
     def unload(self):
         if self.is_loaded:
@@ -111,74 +122,72 @@ class LLMManager:
             del self.processor
             if self.device.type == "cuda":
                 torch.cuda.empty_cache()
-                logger.info(f"GPU memory allocated after unload: {torch.cuda.memory_allocated()}")
+                logger.info(f"GPU memory cleared: {torch.cuda.memory_allocated()} bytes allocated")
             self.is_loaded = False
-            logger.info(f"LLM {self.model_name} unloaded from {self.device}")
+            self.token_cache.clear()
+            logger.info(f"LLM {self.model_name} unloaded")
 
-    async def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> str:
+    async def generate(self, prompt: str, max_tokens: int = settings.max_tokens, temperature: float = 0.7) -> str:
         if not self.is_loaded:
+            logger.warning("LLM not loaded; attempting reload")
             self.load()
+        if not self.is_loaded:
+            raise HTTPException(status_code=503, detail="LLM model unavailable")
 
-        messages_vlm = [
-            {
-                "role": "system",
-                "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Answer questions considering India as base country and Karnataka as base state. Provide a concise response in one sentence maximum."}]
-            },
-            {
-                "role": "user",
-                "content": [{"type": "text", "text": prompt}]
-            }
-        ]
+        cache_key = f"{prompt}:{max_tokens}:{temperature}"
+        if cache_key in self.token_cache:
+            logger.info("Using cached response")
+            return self.token_cache[cache_key]
 
+        future = asyncio.Future()
+        await request_queue.put({"prompt": prompt, "max_tokens": max_tokens, "temperature": temperature, "future": future})
+        response = await future
+        self.token_cache[cache_key] = response
+        logger.info(f"Generated response: {response}")
+        return response
+
+    async def batch_generate(self, prompts: List[Dict]) -> List[str]:
+        messages_batch = [
+            [
+                {"role": "system", "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Answer questions considering India as base country and Karnataka as base state. Provide a concise response in one sentence maximum."}]},
+                {"role": "user", "content": [{"type": "text", "text": prompt["prompt"]}]}
+            ]
+            for prompt in prompts
+        ]
         try:
             inputs_vlm = self.processor.apply_chat_template(
-                messages_vlm,
+                messages_batch,
                 add_generation_prompt=True,
                 tokenize=True,
                 return_dict=True,
-                return_tensors="pt"
+                return_tensors="pt",
+                padding=True
             ).to(self.device, dtype=torch.bfloat16)
+            with autocast(), torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs_vlm,
+                    max_new_tokens=max(prompt["max_tokens"] for prompt in prompts),
+                    do_sample=True,
+                    top_p=0.9,
+                    temperature=max(prompt["temperature"] for prompt in prompts)
+                )
+            responses = [
+                self.processor.decode(output[input_len:], skip_special_tokens=True)
+                for output, input_len in zip(outputs, inputs_vlm["input_ids"].shape[1])
+            ]
+            logger.info(f"Batch generated {len(responses)} responses")
+            return responses
         except Exception as e:
-            logger.error(f"Error in tokenization: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Tokenization failed: {str(e)}")
-
-        input_len = inputs_vlm["input_ids"].shape[-1]
-
-        with torch.inference_mode():
-            generation = self.model.generate(
-                **inputs_vlm,
-                max_new_tokens=max_tokens,
-                do_sample=True,
-                temperature=temperature
-            )
-            generation = generation[0][input_len:]
-
-        response = self.processor.decode(generation, skip_special_tokens=True)
-        logger.info(f"Generated response: {response}")
-        return response
+            logger.error(f"Error in batch generation: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Batch generation failed: {str(e)}")
 
     async def vision_query(self, image: Image.Image, query: str) -> str:
         if not self.is_loaded:
             self.load()
-
         messages_vlm = [
-            {
-                "role": "system",
-                "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Summarize your answer in maximum 1 sentence."}]
-            },
-            {
-                "role": "user",
-                "content": []
-            }
+            {"role": "system", "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Summarize your answer in maximum 1 sentence."}]},
+            {"role": "user", "content": [{"type": "text", "text": query}] + ([{"type": "image", "image": image}] if image and image.size[0] > 0 and image.size[1] > 0 else [])}
         ]
-
-        messages_vlm[1]["content"].append({"type": "text", "text": query})
-        if image and image.size[0] > 0 and image.size[1] > 0:
-            messages_vlm[1]["content"].insert(0, {"type": "image", "image": image})
-            logger.info(f"Received valid image for processing")
-        else:
-            logger.info("No valid image provided, processing text only")
-
         try:
             inputs_vlm = self.processor.apply_chat_template(
                 messages_vlm,
@@ -190,18 +199,10 @@ class LLMManager:
         except Exception as e:
             logger.error(f"Error in apply_chat_template: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Failed to process input: {str(e)}")
-
         input_len = inputs_vlm["input_ids"].shape[-1]
-
         with torch.inference_mode():
-            generation = self.model.generate(
-                **inputs_vlm,
-                max_new_tokens=512,
-                do_sample=True,
-                temperature=0.7
-            )
+            generation = self.model.generate(**inputs_vlm, max_new_tokens=512, do_sample=True, temperature=0.7)
             generation = generation[0][input_len:]
-
         decoded = self.processor.decode(generation, skip_special_tokens=True)
         logger.info(f"Vision query response: {decoded}")
         return decoded
@@ -209,25 +210,10 @@ class LLMManager:
     async def chat_v2(self, image: Image.Image, query: str) -> str:
         if not self.is_loaded:
             self.load()
-
         messages_vlm = [
-            {
-                "role": "system",
-                "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Answer questions considering India as base country and Karnataka as base state."}]
-            },
-            {
-                "role": "user",
-                "content": []
-            }
+            {"role": "system", "content": [{"type": "text", "text": "You are Dhwani, a helpful assistant. Answer questions considering India as base country and Karnataka as base state."}]},
+            {"role": "user", "content": [{"type": "text", "text": query}] + ([{"type": "image", "image": image}] if image and image.size[0] > 0 and image.size[1] > 0 else [])}
         ]
-
-        messages_vlm[1]["content"].append({"type": "text", "text": query})
-        if image and image.size[0] > 0 and image.size[1] > 0:
-            messages_vlm[1]["content"].insert(0, {"type": "image", "image": image})
-            logger.info(f"Received valid image for processing")
-        else:
-            logger.info("No valid image provided, processing text only")
-
         try:
             inputs_vlm = self.processor.apply_chat_template(
                 messages_vlm,
@@ -239,18 +225,10 @@ class LLMManager:
         except Exception as e:
             logger.error(f"Error in apply_chat_template: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Failed to process input: {str(e)}")
-
         input_len = inputs_vlm["input_ids"].shape[-1]
-
         with torch.inference_mode():
-            generation = self.model.generate(
-                **inputs_vlm,
-                max_new_tokens=512,
-                do_sample=True,
-                temperature=0.7
-            )
+            generation = self.model.generate(**inputs_vlm, max_new_tokens=512, do_sample=True, temperature=0.7)
             generation = generation[0][input_len:]
-
         decoded = self.processor.decode(generation, skip_special_tokens=True)
         logger.info(f"Chat_v2 response: {decoded}")
         return decoded
@@ -258,101 +236,42 @@ class LLMManager:
 # TTS Manager
 class TTSManager:
     def __init__(self, device_type=device):
-        self.device_type = device_type
+        self.device_type = torch.device(device_type)
         self.model = None
         self.repo_id = "ai4bharat/IndicF5"
+        self.load()
 
     def load(self):
         if not self.model:
-            logger.info("Loading TTS model IndicF5...")
-            self.model = AutoModel.from_pretrained(
-                self.repo_id,
-                trust_remote_code=True
-            )
-            self.model = self.model.to(self.device_type)
-            logger.info("TTS model IndicF5 loaded")
+            logger.info(f"Loading TTS model {self.repo_id} on {self.device_type}...")
+            self.model = AutoModel.from_pretrained(self.repo_id, trust_remote_code=True).to(self.device_type)
+            logger.info("TTS model loaded")
+
+    def unload(self):
+        if self.model:
+            del self.model
+            if self.device_type.type == "cuda":
+                torch.cuda.empty_cache()
+                logger.info(f"TTS GPU memory cleared: {torch.cuda.memory_allocated()} bytes allocated")
+            self.model = None
+            logger.info("TTS model unloaded")
 
     def synthesize(self, text, ref_audio_path, ref_text):
         if not self.model:
             raise ValueError("TTS model not loaded")
-        return self.model(text, ref_audio_path=ref_audio_path, ref_text=ref_text)
-
-# TTS Constants
-EXAMPLES = [
-    {
-        "audio_name": "KAN_F (Happy)",
-        "audio_url": "https://github.com/AI4Bharat/IndicF5/raw/refs/heads/main/prompts/KAN_F_HAPPY_00001.wav",
-        "ref_text": "ನಮ್‌ ಫ್ರಿಜ್ಜಲ್ಲಿ  ಕೂಲಿಂಗ್‌ ಸಮಸ್ಯೆ ಆಗಿ ನಾನ್‌ ಭಾಳ ದಿನದಿಂದ ಒದ್ದಾಡ್ತಿದ್ದೆ, ಆದ್ರೆ ಅದ್ನೀಗ ಮೆಕಾನಿಕ್ ಆಗಿರೋ ನಿಮ್‌ ಸಹಾಯ್ದಿಂದ ಬಗೆಹರಿಸ್ಕೋಬೋದು ಅಂತಾಗಿ ನಿರಾಳ ಆಯ್ತು ನಂಗೆ.",
-        "synth_text": "ಚೆನ್ನೈನ ಶೇರ್ ಆಟೋ ಪ್ರಯಾಣಿಕರ ನಡುವೆ ಆಹಾರವನ್ನು ಹಂಚಿಕೊಂಡು ತಿನ್ನುವುದು ನನಗೆ ಮನಸ್ಸಿಗೆ ತುಂಬಾ ಒಳ್ಳೆಯದೆನಿ��ುವ ವಿಷಯ."
-    },
-]
-
-# Pydantic models for TTS
-class SynthesizeRequest(BaseModel):
-    text: str
-    ref_audio_name: str
-    ref_text: str = None
-
-class KannadaSynthesizeRequest(BaseModel):
-    text: str
-
-# TTS Functions
-def load_audio_from_url(url: str):
-    response = requests.get(url)
-    if response.status_code == 200:
-        audio_data, sample_rate = sf.read(io.BytesIO(response.content))
-        return sample_rate, audio_data
-    raise HTTPException(status_code=500, detail="Failed to load reference audio from URL.")
-
-def synthesize_speech(tts_manager: TTSManager, text: str, ref_audio_name: str, ref_text: str):
-    ref_audio_url = None
-    for example in EXAMPLES:
-        if example["audio_name"] == ref_audio_name:
-            ref_audio_url = example["audio_url"]
-            if not ref_text:
-                ref_text = example["ref_text"]
-            break
-    
-    if not ref_audio_url:
-        raise HTTPException(status_code=400, detail="Invalid reference audio name.")
-    if not text.strip():
-        raise HTTPException(status_code=400, detail="Text to synthesize cannot be empty.")
-    if not ref_text or not ref_text.strip():
-        raise HTTPException(status_code=400, detail="Reference text cannot be empty.")
-
-    sample_rate, audio_data = load_audio_from_url(ref_audio_url)
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
-        sf.write(temp_audio.name, audio_data, samplerate=sample_rate, format='WAV')
-        temp_audio.flush()
-        audio = tts_manager.synthesize(text, ref_audio_path=temp_audio.name, ref_text=ref_text)
-
-    if audio.dtype == np.int16:
-        audio = audio.astype(np.float32) / 32768.0
-    buffer = io.BytesIO()
-    sf.write(buffer, audio, 24000, format='WAV')
-    buffer.seek(0)
-    return buffer
-
-# Supported languages
-SUPPORTED_LANGUAGES = {
-    "asm_Beng", "kas_Arab", "pan_Guru", "ben_Beng", "kas_Deva", "san_Deva",
-    "brx_Deva", "mai_Deva", "sat_Olck", "doi_Deva", "mal_Mlym", "snd_Arab",
-    "eng_Latn", "mar_Deva", "snd_Deva", "gom_Deva", "mni_Beng", "tam_Taml",
-    "guj_Gujr", "mni_Mtei", "tel_Telu", "hin_Deva", "npi_Deva", "urd_Arab",
-    "kan_Knda", "ory_Orya",
-    "deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn",
-    "por_Latn", "rus_Cyrl", "pol_Latn"
-}
+        with autocast():
+            return self.model(text, ref_audio_path=ref_audio_path, ref_text=ref_text)
 
 # Translation Manager
 class TranslateManager:
     def __init__(self, src_lang, tgt_lang, device_type=device, use_distilled=True):
-        self.device_type = device_type
+        self.device_type = torch.device(device_type)
         self.tokenizer = None
         self.model = None
         self.src_lang = src_lang
         self.tgt_lang = tgt_lang
         self.use_distilled = use_distilled
+        self.load()
 
     def load(self):
         if not self.tokenizer or not self.model:
@@ -364,21 +283,17 @@ class TranslateManager:
                 model_name = "ai4bharat/indictrans2-indic-indic-dist-320M" if self.use_distilled else "ai4bharat/indictrans2-indic-indic-1B"
             else:
                 raise ValueError("Invalid language combination")
-
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                model_name,
-                trust_remote_code=True
-            )
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
             self.model = AutoModelForSeq2SeqLM.from_pretrained(
                 model_name,
                 trust_remote_code=True,
                 torch_dtype=torch.float16,
                 attn_implementation="flash_attention_2"
-            )
-            self.model = self.model.to(self.device_type)
+            ).to(self.device_type)
             self.model = torch.compile(self.model, mode="reduce-overhead")
             logger.info(f"Translation model {model_name} loaded")
 
+# Model Manager
 class ModelManager:
     def __init__(self, device_type=device, use_distilled=True, is_lazy_loading=False):
         self.models = {}
@@ -389,18 +304,14 @@ class ModelManager:
     def load_model(self, src_lang, tgt_lang, key):
         logger.info(f"Loading translation model for {src_lang} -> {tgt_lang}")
         translate_manager = TranslateManager(src_lang, tgt_lang, self.device_type, self.use_distilled)
-        translate_manager.load()
         self.models[key] = translate_manager
         logger.info(f"Loaded translation model for {key}")
 
     def get_model(self, src_lang, tgt_lang):
         key = self._get_model_key(src_lang, tgt_lang)
-        if key not in self.models:
-            if self.is_lazy_loading:
-                self.load_model(src_lang, tgt_lang, key)
-            else:
-                raise ValueError(f"Model for {key} is not preloaded and lazy loading is disabled.")
-        return self.models.get(key)
+        if key not in self.models and self.is_lazy_loading:
+            self.load_model(src_lang, tgt_lang, key)
+        return self.models.get(key) or (self.load_model(src_lang, tgt_lang, key) or self.models[key])
 
     def _get_model_key(self, src_lang, tgt_lang):
         if src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
@@ -413,21 +324,30 @@ class ModelManager:
 
 # ASR Manager
 class ASRModelManager:
-    def __init__(self, device_type="cuda"):
-        self.device_type = device_type
+    def __init__(self, device_type=device):
+        self.device_type = torch.device(device_type)
         self.model = None
         self.model_language = {"kannada": "kn"}
+        self.load()
 
     def load(self):
         if not self.model:
-            logger.info("Loading ASR model...")
+            logger.info(f"Loading ASR model on {self.device_type}...")
             self.model = AutoModel.from_pretrained(
                 "ai4bharat/indic-conformer-600m-multilingual",
                 trust_remote_code=True
-            )
-            self.model = self.model.to(self.device_type)
+            ).to(self.device_type)
             logger.info("ASR model loaded")
 
+    def unload(self):
+        if self.model:
+            del self.model
+            if self.device_type.type == "cuda":
+                torch.cuda.empty_cache()
+                logger.info(f"ASR GPU memory cleared: {torch.cuda.memory_allocated()} bytes allocated")
+            self.model = None
+            logger.info("ASR model unloaded")
+
 # Global Managers
 llm_manager = LLMManager(settings.llm_model_name)
 model_manager = ModelManager()
@@ -435,7 +355,31 @@ asr_manager = ASRModelManager()
 tts_manager = TTSManager()
 ip = IndicProcessor(inference=True)
 
+# TTS Constants
+EXAMPLES = [
+    {
+        "audio_name": "KAN_F (Happy)",
+        "audio_url": "https://github.com/AI4Bharat/IndicF5/raw/refs/heads/main/prompts/KAN_F_HAPPY_00001.wav",
+        "ref_text": "ನಮ್‌ ಫ್ರಿಜ್ಜಲ್ಲಿ ಕೂಲಿಂಗ್‌ ಸಮಸ್ಯೆ ಆಗಿ ನಾನ್‌ ಭಾಳ ದಿನದಿಂದ ಒದ್ದಾಡ್ತಿದ್ದೆ, ಆದ್ರೆ ಅದ್ನೀಗ ಮೆಕಾನಿಕ್ ಆಗಿರೋ ನಿಮ್‌ ಸಹಾಯ್ದಿಂದ ಬಗೆಹರಿಸ್ಕೋಬೋದು ಅಂತಾಗಿ ನಿರಾಳ ಆಯ್ತು ನಂಗೆ।",
+        "synth_text": "ಚೆನ್ನೈನ ಶೇರ್ ಆಟೋ ಪ್ರಯಾಣಿಕರ ನಡುವೆ ಆಹಾರವನ್ನು ಹಂಚಿಕೊಂಡು ತಿನ್ನುವುದು ನನಗೆ ಮನಸ್ಸಿಗೆ ತುಂಬಾ ಒಳ್ಳೆಯದೆನಿಸುವ ವಿಷಯ."
+    },
+]
+
 # Pydantic Models
+class SynthesizeRequest(BaseModel):
+    text: str
+    ref_audio_name: str
+    ref_text: str = None
+
+class KannadaSynthesizeRequest(BaseModel):
+    text: str
+
+    @field_validator("text")
+    def text_must_be_valid(cls, v):
+        if len(v) > 500:
+            raise ValueError("Text cannot exceed 500 characters")
+        return v.strip()
+
 class ChatRequest(BaseModel):
     prompt: str
     src_lang: str = "kan_Knda"
@@ -453,7 +397,6 @@ class ChatRequest(BaseModel):
             raise ValueError(f"Unsupported language code: {v}. Supported codes: {', '.join(SUPPORTED_LANGUAGES)}")
         return v
 
-
 class ChatResponse(BaseModel):
     response: str
 
@@ -468,71 +411,149 @@ class TranscriptionResponse(BaseModel):
 class TranslationResponse(BaseModel):
     translations: List[str]
 
+# TTS Functions
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10))
+def load_audio_from_url(url: str):
+    response = requests.get(url)
+    if response.status_code == 200:
+        audio_data, sample_rate = sf.read(io.BytesIO(response.content))
+        return sample_rate, audio_data
+    raise HTTPException(status_code=500, detail="Failed to load reference audio from URL after retries")
+
+async def synthesize_speech(tts_manager: TTSManager, text: str, ref_audio_name: str, ref_text: str) -> io.BytesIO:
+    async with request_queue:
+        ref_audio_url = next((ex["audio_url"] for ex in EXAMPLES if ex["audio_name"] == ref_audio_name), None)
+        if not ref_audio_url:
+            raise HTTPException(status_code=400, detail="Invalid reference audio name.")
+        if not text.strip() or not ref_text.strip():
+            raise HTTPException(status_code=400, detail="Text or reference text cannot be empty.")
+
+        logger.info(f"Synthesizing speech for text: {text[:50]}... with ref_audio: {ref_audio_name}")
+        loop = asyncio.get_running_loop()
+        sample_rate, audio_data = await loop.run_in_executor(None, load_audio_from_url, ref_audio_url)
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as temp_audio:
+            await loop.run_in_executor(None, sf.write, temp_audio.name, audio_data, sample_rate, "WAV")
+            temp_audio.flush()
+            audio = tts_manager.synthesize(text, temp_audio.name, ref_text)
+
+        buffer = io.BytesIO()
+        await loop.run_in_executor(None, sf.write, buffer, audio.astype(np.float32) / 32768.0 if audio.dtype == np.int16 else audio, 24000, "WAV")
+        buffer.seek(0)
+        logger.info("Speech synthesis completed")
+        return buffer
+
+# Supported Languages
+SUPPORTED_LANGUAGES = {
+    "asm_Beng", "kas_Arab", "pan_Guru", "ben_Beng", "kas_Deva", "san_Deva",
+    "brx_Deva", "mai_Deva", "sat_Olck", "doi_Deva", "mal_Mlym", "snd_Arab",
+    "eng_Latn", "mar_Deva", "snd_Deva", "gom_Deva", "mni_Beng", "tam_Taml",
+    "guj_Gujr", "mni_Mtei", "tel_Telu", "hin_Deva", "npi_Deva", "urd_Arab",
+    "kan_Knda", "ory_Orya",
+    "deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn",
+    "por_Latn", "rus_Cyrl", "pol_Latn"
+}
+
 # Dependency
 def get_translate_manager(src_lang: str, tgt_lang: str) -> TranslateManager:
     return model_manager.get_model(src_lang, tgt_lang)
 
+# Translation Function
+async def perform_internal_translation(sentences: List[str], src_lang: str, tgt_lang: str) -> List[str]:
+    try:
+        translate_manager = model_manager.get_model(src_lang, tgt_lang)
+    except ValueError as e:
+        logger.info(f"Model not preloaded: {str(e)}, loading now...")
+        key = model_manager._get_model_key(src_lang, tgt_lang)
+        model_manager.load_model(src_lang, tgt_lang, key)
+        translate_manager = model_manager.get_model(src_lang, tgt_lang)
+
+    batch = ip.preprocess_batch(sentences, src_lang=src_lang, tgt_lang=tgt_lang)
+    inputs = translate_manager.tokenizer(batch, truncation=True, padding="longest", return_tensors="pt", return_attention_mask=True).to(translate_manager.device_type)
+    with torch.no_grad(), autocast():
+        generated_tokens = translate_manager.model.generate(**inputs, use_cache=True, min_length=0, max_length=256, num_beams=5, num_return_sequences=1)
+    with translate_manager.tokenizer.as_target_tokenizer():
+        generated_tokens = translate_manager.tokenizer.batch_decode(generated_tokens.detach().cpu().tolist(), skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    return ip.postprocess_batch(generated_tokens, lang=tgt_lang)
+
 # Lifespan Event Handler
 translation_configs = []
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     def load_all_models():
-        try:
-            # Load LLM model
-            logger.info("Loading LLM model...")
-            llm_manager.load()
-            logger.info("LLM model loaded successfully")
-
-            # Load TTS model
-            logger.info("Loading TTS model...")
-            tts_manager.load()
-            logger.info("TTS model loaded successfully")
-
-            # Load ASR model
-            logger.info("Loading ASR model...")
-            asr_manager.load()
-            logger.info("ASR model loaded successfully")
-
-            # Load translation models
-            translation_tasks = [
-                ('eng_Latn', 'kan_Knda', 'eng_indic'),
-                ('kan_Knda', 'eng_Latn', 'indic_eng'),
-                ('kan_Knda', 'hin_Deva', 'indic_indic'),
-            ]
-            
-            for config in translation_configs:
-                src_lang = config["src_lang"]
-                tgt_lang = config["tgt_lang"]
-                key = model_manager._get_model_key(src_lang, tgt_lang)
-                translation_tasks.append((src_lang, tgt_lang, key))
-
-            for src_lang, tgt_lang, key in translation_tasks:
-                logger.info(f"Loading translation model for {src_lang} -> {tgt_lang}...")
-                model_manager.load_model(src_lang, tgt_lang, key)
-                logger.info(f"Translation model for {key} loaded successfully")
-
-            logger.info("All models loaded successfully")
-        except Exception as e:
-            logger.error(f"Error loading models: {str(e)}")
-            raise
+        logger.info("Loading LLM model...")
+        llm_manager.load()
+        logger.info("Loading TTS model...")
+        tts_manager.load()
+        logger.info("Loading ASR model...")
+        asr_manager.load()
+        translation_tasks = [
+            ('eng_Latn', 'kan_Knda', 'eng_indic'),
+            ('kan_Knda', 'eng_Latn', 'indic_eng'),
+            ('kan_Knda', 'hin_Deva', 'indic_indic'),
+        ]
+        for config in translation_configs:
+            src_lang = config["src_lang"]
+            tgt_lang = config["tgt_lang"]
+            key = model_manager._get_model_key(src_lang, tgt_lang)
+            translation_tasks.append((src_lang, tgt_lang, key))
+        for src_lang, tgt_lang, key in translation_tasks:
+            logger.info(f"Loading translation model for {src_lang} -> {tgt_lang}...")
+            model_manager.load_model(src_lang, tgt_lang, key)
+        logger.info("All models loaded successfully")
 
-    logger.info("Starting sequential model loading...")
+    logger.info("Starting server with preloaded models...")
     load_all_models()
+    batch_task = asyncio.create_task(batch_worker())
     yield
+    batch_task.cancel()
     llm_manager.unload()
-    logger.info("Server shutdown complete")
+    tts_manager.unload()
+    asr_manager.unload()
+    for model in model_manager.models.values():
+        model.unload()
+    logger.info("Server shutdown complete; all models unloaded")
+
+# Batch Worker
+async def batch_worker():
+    while True:
+        batch = []
+        last_request_time = time()
+        try:
+            while len(batch) < 4:
+                try:
+                    request = await asyncio.wait_for(request_queue.get(), timeout=1.0)
+                    batch.append(request)
+                    current_time = time()
+                    if current_time - last_request_time > 1.0 and batch:
+                        break
+                    last_request_time = current_time
+                except asyncio.TimeoutError:
+                    if batch:
+                        break
+                    continue
+            if batch:
+                start_time = time()
+                responses = await llm_manager.batch_generate(batch)
+                duration = time() - start_time
+                logger.info(f"Batch of {len(batch)} requests processed in {duration:.3f} seconds")
+                for request, response in zip(batch, responses):
+                    request["future"].set_result(response)
+        except Exception as e:
+            logger.error(f"Batch worker error: {str(e)}")
+            for request in batch:
+                request["future"].set_exception(e)
 
 # FastAPI App
 app = FastAPI(
-    title="Dhwani API",
-    description="AI Chat API supporting Indian languages",
+    title="Optimized Dhwani API",
+    description="AI Chat API supporting Indian languages with performance enhancements",
     version="1.0.0",
     redirect_slashes=False,
     lifespan=lifespan
 )
 
-# Add CORS Middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -541,13 +562,11 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Add Timing Middleware
 @app.middleware("http")
 async def add_request_timing(request: Request, call_next):
     start_time = time()
     response = await call_next(request)
-    end_time = time()
-    duration = end_time - start_time
+    duration = time() - start_time
     logger.info(f"Request to {request.url.path} took {duration:.3f} seconds")
     response.headers["X-Response-Time"] = f"{duration:.3f}"
     return response
@@ -555,7 +574,7 @@ async def add_request_timing(request: Request, call_next):
 limiter = Limiter(key_func=get_remote_address)
 app.state.limiter = limiter
 
-# API Endpoints
+# Endpoints
 @app.post("/audio/speech", response_class=StreamingResponse)
 async def synthesize_kannada(request: KannadaSynthesizeRequest):
     if not tts_manager.model:
@@ -563,14 +582,7 @@ async def synthesize_kannada(request: KannadaSynthesizeRequest):
     kannada_example = next(ex for ex in EXAMPLES if ex["audio_name"] == "KAN_F (Happy)")
     if not request.text.strip():
         raise HTTPException(status_code=400, detail="Text to synthesize cannot be empty.")
-    
-    audio_buffer = synthesize_speech(
-        tts_manager,
-        text=request.text,
-        ref_audio_name="KAN_F (Happy)",
-        ref_text=kannada_example["ref_text"]
-    )
-    
+    audio_buffer = await synthesize_speech(tts_manager, request.text, "KAN_F (Happy)", kannada_example["ref_text"])
     return StreamingResponse(
         audio_buffer,
         media_type="audio/wav",
@@ -579,61 +591,69 @@ async def synthesize_kannada(request: KannadaSynthesizeRequest):
 
 @app.post("/translate", response_model=TranslationResponse)
 async def translate(request: TranslationRequest, translate_manager: TranslateManager = Depends(get_translate_manager)):
-    input_sentences = request.sentences
-    src_lang = request.src_lang
-    tgt_lang = request.tgt_lang
-
-    if not input_sentences:
+    if not request.sentences:
         raise HTTPException(status_code=400, detail="Input sentences are required")
-
-    batch = ip.preprocess_batch(input_sentences, src_lang=src_lang, tgt_lang=tgt_lang)
-    inputs = translate_manager.tokenizer(
-        batch,
-        truncation=True,
-        padding="longest",
-        return_tensors="pt",
-        return_attention_mask=True,
-    ).to(translate_manager.device_type)
-
-    with torch.no_grad():
-        generated_tokens = translate_manager.model.generate(
-            **inputs,
-            use_cache=True,
-            min_length=0,
-            max_length=256,
-            num_beams=5,
-            num_return_sequences=1,
-        )
-
+    batch = ip.preprocess_batch(request.sentences, src_lang=request.src_lang, tgt_lang=request.tgt_lang)
+    inputs = translate_manager.tokenizer(batch, truncation=True, padding="longest", return_tensors="pt", return_attention_mask=True).to(translate_manager.device_type)
+    with torch.no_grad(), autocast():
+        generated_tokens = translate_manager.model.generate(**inputs, use_cache=True, min_length=0, max_length=256, num_beams=5, num_return_sequences=1)
     with translate_manager.tokenizer.as_target_tokenizer():
-        generated_tokens = translate_manager.tokenizer.batch_decode(
-            generated_tokens.detach().cpu().tolist(),
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True,
-        )
-
-    translations = ip.postprocess_batch(generated_tokens, lang=tgt_lang)
+        generated_tokens = translate_manager.tokenizer.batch_decode(generated_tokens.detach().cpu().tolist(), skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    translations = ip.postprocess_batch(generated_tokens, lang=request.tgt_lang)
     return TranslationResponse(translations=translations)
 
-async def perform_internal_translation(sentences: List[str], src_lang: str, tgt_lang: str) -> List[str]:
-    try:
-        translate_manager = model_manager.get_model(src_lang, tgt_lang)
-    except ValueError as e:
-        logger.info(f"Model not preloaded: {str(e)}, loading now...")
-        key = model_manager._get_model_key(src_lang, tgt_lang)
-        model_manager.load_model(src_lang, tgt_lang, key)
-        translate_manager = model_manager.get_model(src_lang, tgt_lang)
-    
-    if not translate_manager.model:
-        translate_manager.load()
-    
-    request = TranslationRequest(sentences=sentences, src_lang=src_lang, tgt_lang=tgt_lang)
-    response = await translate(request, translate_manager)
-    return response.translations
-
 @app.get("/v1/health")
 async def health_check():
-    return {"status": "healthy", "model": settings.llm_model_name}
+    memory_usage = torch.cuda.memory_allocated() / (24 * 1024**3) if cuda_available else 0
+    if memory_usage > 0.9:
+        logger.warning("GPU memory usage exceeds 90%; consider unloading models")
+    llm_status = "unhealthy"
+    llm_latency = None
+    if llm_manager.is_loaded:
+        start = time()
+        try:
+            llm_test = await llm_manager.generate("What is the capital of Karnataka?", max_tokens=10)
+            llm_latency = time() - start
+            llm_status = "healthy" if llm_test else "unhealthy"
+        except Exception as e:
+            logger.error(f"LLM health check failed: {str(e)}")
+    tts_status = "unhealthy"
+    tts_latency = None
+    if tts_manager.model:
+        start = time()
+        try:
+            audio_buffer = await synthesize_speech(tts_manager, "Test", "KAN_F (Happy)", EXAMPLES[0]["ref_text"])
+            tts_latency = time() - start
+            tts_status = "healthy" if audio_buffer else "unhealthy"
+        except Exception as e:
+            logger.error(f"TTS health check failed: {str(e)}")
+    asr_status = "unhealthy"
+    asr_latency = None
+    if asr_manager.model:
+        start = time()
+        try:
+            dummy_audio = np.zeros(16000, dtype=np.float32)
+            wav = torch.tensor(dummy_audio).unsqueeze(0).to(device)
+            with autocast(), torch.no_grad():
+                asr_test = asr_manager.model(wav, asr_manager.model_language["kannada"], "rnnt")
+            asr_latency = time() - start
+            asr_status = "healthy" if asr_test else "unhealthy"
+        except Exception as e:
+            logger.error(f"ASR health check failed: {str(e)}")
+    status = {
+        "status": "healthy" if llm_status == "healthy" and tts_status == "healthy" and asr_status == "healthy" else "degraded",
+        "model": settings.llm_model_name,
+        "llm_status": llm_status,
+        "llm_latency": f"{llm_latency:.3f}s" if llm_latency else "N/A",
+        "tts_status": tts_status,
+        "tts_latency": f"{tts_latency:.3f}s" if tts_latency else "N/A",
+        "asr_status": asr_status,
+        "asr_latency": f"{asr_latency:.3f}s" if asr_latency else "N/A",
+        "translation_models": list(model_manager.models.keys()),
+        "gpu_memory_usage": f"{memory_usage:.2%}"
+    }
+    logger.info("Health check completed")
+    return status
 
 @app.get("/")
 async def home():
@@ -644,6 +664,10 @@ async def unload_all_models():
     try:
         logger.info("Starting to unload all models...")
         llm_manager.unload()
+        tts_manager.unload()
+        asr_manager.unload()
+        for model in model_manager.models.values():
+            model.unload()
         logger.info("All models unloaded successfully")
         return {"status": "success", "message": "All models unloaded"}
     except Exception as e:
@@ -655,6 +679,15 @@ async def load_all_models():
     try:
         logger.info("Starting to load all models...")
         llm_manager.load()
+        tts_manager.load()
+        asr_manager.load()
+        for src_lang, tgt_lang, key in [
+            ('eng_Latn', 'kan_Knda', 'eng_indic'),
+            ('kan_Knda', 'eng_Latn', 'indic_eng'),
+            ('kan_Knda', 'hin_Deva', 'indic_indic'),
+        ]:
+            if key not in model_manager.models:
+                model_manager.load_model(src_lang, tgt_lang, key)
         logger.info("All models loaded successfully")
         return {"status": "success", "message": "All models loaded"}
     except Exception as e:
@@ -665,11 +698,7 @@ async def load_all_models():
 async def translate_endpoint(request: TranslationRequest):
     logger.info(f"Received translation request: {request.dict()}")
     try:
-        translations = await perform_internal_translation(
-            sentences=request.sentences,
-            src_lang=request.src_lang,
-            tgt_lang=request.tgt_lang
-        )
+        translations = await perform_internal_translation(request.sentences, request.src_lang, request.tgt_lang)
         logger.info(f"Translation successful: {translations}")
         return TranslationResponse(translations=translations)
     except Exception as e:
@@ -679,44 +708,32 @@ async def translate_endpoint(request: TranslationRequest):
 @app.post("/v1/chat", response_model=ChatResponse)
 @limiter.limit(settings.chat_rate_limit)
 async def chat(request: Request, chat_request: ChatRequest):
-    if not chat_request.prompt:
-        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-    logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, tgt_lang: {chat_request.tgt_lang}")
-    
-    EUROPEAN_LANGUAGES = {"deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn", "por_Latn", "rus_Cyrl", "pol_Latn"}
-    
-    try:
-        if chat_request.src_lang != "eng_Latn" and chat_request.src_lang not in EUROPEAN_LANGUAGES:
-            translated_prompt = await perform_internal_translation(
-                sentences=[chat_request.prompt],
-                src_lang=chat_request.src_lang,
-                tgt_lang="eng_Latn"
-            )
-            prompt_to_process = translated_prompt[0]
-            logger.info(f"Translated prompt to English: {prompt_to_process}")
-        else:
-            prompt_to_process = chat_request.prompt
-            logger.info("Prompt in English or European language, no translation needed")
-
-        response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-        logger.info(f"Generated response: {response}")
-
-        if chat_request.tgt_lang != "eng_Latn" and chat_request.tgt_lang not in EUROPEAN_LANGUAGES:
-            translated_response = await perform_internal_translation(
-                sentences=[response],
-                src_lang="eng_Latn",
-                tgt_lang=chat_request.tgt_lang
-            )
-            final_response = translated_response[0]
-            logger.info(f"Translated response to {chat_request.tgt_lang}: {final_response}")
-        else:
-            final_response = response
-            logger.info(f"Response in {chat_request.tgt_lang}, no translation needed")
-
-        return ChatResponse(response=final_response)
-    except Exception as e:
-        logger.error(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    async with request_queue:
+        if not chat_request.prompt:
+            raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, tgt_lang: {chat_request.tgt_lang}")
+        EUROPEAN_LANGUAGES = {"deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn", "por_Latn", "rus_Cyrl", "pol_Latn"}
+        try:
+            if chat_request.src_lang != "eng_Latn" and chat_request.src_lang not in EUROPEAN_LANGUAGES:
+                translated_prompt = await perform_internal_translation([chat_request.prompt], chat_request.src_lang, "eng_Latn")
+                prompt_to_process = translated_prompt[0]
+                logger.info(f"Translated prompt to English: {prompt_to_process}")
+            else:
+                prompt_to_process = chat_request.prompt
+                logger.info("Prompt in English or European language, no translation needed")
+            response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
+            logger.info(f"Generated English response: {response}")
+            if chat_request.tgt_lang != "eng_Latn" and chat_request.tgt_lang not in EUROPEAN_LANGUAGES:
+                translated_response = await perform_internal_translation([response], "eng_Latn", chat_request.tgt_lang)
+                final_response = translated_response[0]
+                logger.info(f"Translated response to {chat_request.tgt_lang}: {final_response}")
+            else:
+                final_response = response
+                logger.info(f"Response in {chat_request.tgt_lang}, no translation needed")
+            return ChatResponse(response=final_response)
+        except Exception as e:
+            logger.error(f"Error processing request: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 @app.post("/v1/visual_query/")
 async def visual_query(
@@ -725,42 +742,31 @@ async def visual_query(
     src_lang: str = Query("kan_Knda", enum=list(SUPPORTED_LANGUAGES)),
     tgt_lang: str = Query("kan_Knda", enum=list(SUPPORTED_LANGUAGES)),
 ):
-    try:
-        image = Image.open(file.file)
-        if image.size == (0, 0):
-            raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
-
-        if src_lang != "eng_Latn":
-            translated_query = await perform_internal_translation(
-                sentences=[query],
-                src_lang=src_lang,
-                tgt_lang="eng_Latn"
-            )
-            query_to_process = translated_query[0]
-            logger.info(f"Translated query to English: {query_to_process}")
-        else:
-            query_to_process = query
-            logger.info("Query already in English, no translation needed")
-
-        answer = await llm_manager.vision_query(image, query_to_process)
-        logger.info(f"Generated English answer: {answer}")
-
-        if tgt_lang != "eng_Latn":
-            translated_answer = await perform_internal_translation(
-                sentences=[answer],
-                src_lang="eng_Latn",
-                tgt_lang=tgt_lang
-            )
-            final_answer = translated_answer[0]
-            logger.info(f"Translated answer to {tgt_lang}: {final_answer}")
-        else:
-            final_answer = answer
-            logger.info("Answer kept in English, no translation needed")
-
-        return {"answer": final_answer}
-    except Exception as e:
-        logger.error(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+    async with request_queue:
+        try:
+            image = Image.open(file.file)
+            if image.size == (0, 0):
+                raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
+            if src_lang != "eng_Latn":
+                translated_query = await perform_internal_translation([query], src_lang, "eng_Latn")
+                query_to_process = translated_query[0]
+                logger.info(f"Translated query to English: {query_to_process}")
+            else:
+                query_to_process = query
+                logger.info("Query already in English, no translation needed")
+            answer = await llm_manager.vision_query(image, query_to_process)
+            logger.info(f"Generated English answer: {answer}")
+            if tgt_lang != "eng_Latn":
+                translated_answer = await perform_internal_translation([answer], "eng_Latn", tgt_lang)
+                final_answer = translated_answer[0]
+                logger.info(f"Translated answer to {tgt_lang}: {final_answer}")
+            else:
+                final_answer = answer
+                logger.info("Answer kept in English, no translation needed")
+            return {"answer": final_answer}
+        except Exception as e:
+            logger.error(f"Error processing request: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 @app.post("/v1/chat_v2", response_model=ChatResponse)
 @limiter.limit(settings.chat_rate_limit)
@@ -771,95 +777,70 @@ async def chat_v2(
     src_lang: str = Form("kan_Knda"),
     tgt_lang: str = Form("kan_Knda"),
 ):
-    if not prompt:
-        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-    if src_lang not in SUPPORTED_LANGUAGES or tgt_lang not in SUPPORTED_LANGUAGES:
-        raise HTTPException(status_code=400, detail=f"Unsupported language code. Supported codes: {', '.join(SUPPORTED_LANGUAGES)}")
-
-    logger.info(f"Received prompt: {prompt}, src_lang: {src_lang}, tgt_lang: {tgt_lang}, Image provided: {image is not None}")
-
-    try:
-        if image:
-            image_data = await image.read()
-            if not image_data:
-                raise HTTPException(status_code=400, detail="Uploaded image is empty")
-            img = Image.open(io.BytesIO(image_data))
-
-            if src_lang != "eng_Latn":
-                translated_prompt = await perform_internal_translation(
-                    sentences=[prompt],
-                    src_lang=src_lang,
-                    tgt_lang="eng_Latn"
-                )
-                prompt_to_process = translated_prompt[0]
-                logger.info(f"Translated prompt to English: {prompt_to_process}")
-            else:
-                prompt_to_process = prompt
-                logger.info("Prompt already in English, no translation needed")
-
-            decoded = await llm_manager.chat_v2(img, prompt_to_process)
-            logger.info(f"Generated English response: {decoded}")
-
-            if tgt_lang != "eng_Latn":
-                translated_response = await perform_internal_translation(
-                    sentences=[decoded],
-                    src_lang="eng_Latn",
-                    tgt_lang=tgt_lang
-                )
-                final_response = translated_response[0]
-                logger.info(f"Translated response to {tgt_lang}: {final_response}")
-            else:
-                final_response = decoded
-                logger.info("Response kept in English, no translation needed")
-        else:
-            if src_lang != "eng_Latn":
-                translated_prompt = await perform_internal_translation(
-                    sentences=[prompt],
-                    src_lang=src_lang,
-                    tgt_lang="eng_Latn"
-                )
-                prompt_to_process = translated_prompt[0]
-                logger.info(f"Translated prompt to English: {prompt_to_process}")
-            else:
-                prompt_to_process = prompt
-                logger.info("Prompt already in English, no translation needed")
-
-            decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-            logger.info(f"Generated English response: {decoded}")
-
-            if tgt_lang != "eng_Latn":
-                translated_response = await perform_internal_translation(
-                    sentences=[decoded],
-                    src_lang="eng_Latn",
-                    tgt_lang=tgt_lang
-                )
-                final_response = translated_response[0]
-                logger.info(f"Translated response to {tgt_lang}: {final_response}")
+    async with request_queue:
+        if not prompt:
+            raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        if src_lang not in SUPPORTED_LANGUAGES or tgt_lang not in SUPPORTED_LANGUAGES:
+            raise HTTPException(status_code=400, detail=f"Unsupported language code. Supported codes: {', '.join(SUPPORTED_LANGUAGES)}")
+        logger.info(f"Received prompt: {prompt}, src_lang: {src_lang}, tgt_lang: {tgt_lang}, Image provided: {image is not None}")
+        try:
+            if image:
+                image_data = await image.read()
+                if not image_data:
+                    raise HTTPException(status_code=400, detail="Uploaded image is empty")
+                img = Image.open(io.BytesIO(image_data))
+                if src_lang != "eng_Latn":
+                    translated_prompt = await perform_internal_translation([prompt], src_lang, "eng_Latn")
+                    prompt_to_process = translated_prompt[0]
+                    logger.info(f"Translated prompt to English: {prompt_to_process}")
+                else:
+                    prompt_to_process = prompt
+                decoded = await llm_manager.chat_v2(img, prompt_to_process)
+                logger.info(f"Generated English response: {decoded}")
+                if tgt_lang != "eng_Latn":
+                    translated_response = await perform_internal_translation([decoded], "eng_Latn", tgt_lang)
+                    final_response = translated_response[0]
+                    logger.info(f"Translated response to {tgt_lang}: {final_response}")
+                else:
+                    final_response = decoded
             else:
-                final_response = decoded
-                logger.info("Response kept in English, no translation needed")
-
-        return ChatResponse(response=final_response)
-    except Exception as e:
-        logger.error(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+                if src_lang != "eng_Latn":
+                    translated_prompt = await perform_internal_translation([prompt], src_lang, "eng_Latn")
+                    prompt_to_process = translated_prompt[0]
+                    logger.info(f"Translated prompt to English: {prompt_to_process}")
+                else:
+                    prompt_to_process = prompt
+                decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
+                logger.info(f"Generated English response: {decoded}")
+                if tgt_lang != "eng_Latn":
+                    translated_response = await perform_internal_translation([decoded], "eng_Latn", tgt_lang)
+                    final_response = translated_response[0]
+                    logger.info(f"Translated response to {tgt_lang}: {final_response}")
+                else:
+                    final_response = decoded
+            return ChatResponse(response=final_response)
+        except Exception as e:
+            logger.error(f"Error processing request: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
 
 @app.post("/transcribe/", response_model=TranscriptionResponse)
 async def transcribe_audio(file: UploadFile = File(...), language: str = Query(..., enum=list(asr_manager.model_language.keys()))):
-    if not asr_manager.model:
-        raise HTTPException(status_code=503, detail="ASR model not loaded")
-    try:
-        wav, sr = torchaudio.load(file.file)
-        wav = torch.mean(wav, dim=0, keepdim=True)
-        target_sample_rate = 16000
-        if sr != target_sample_rate:
-            resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sample_rate)
-            wav = resampler(wav)
-        transcription_rnnt = asr_manager.model(wav, asr_manager.model_language[language], "rnnt")
-        return TranscriptionResponse(text=transcription_rnnt)
-    except Exception as e:
-        logger.error(f"Error in transcription: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
+    async with request_queue:
+        if not asr_manager.model:
+            raise HTTPException(status_code=503, detail="ASR model not loaded")
+        try:
+            wav, sr = torchaudio.load(file.file, backend="cuda" if cuda_available else "cpu")
+            wav = torch.mean(wav, dim=0, keepdim=True).to(device)
+            target_sample_rate = 16000
+            if sr != target_sample_rate:
+                resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sample_rate).to(device)
+                wav = resampler(wav)
+            with autocast(), torch.no_grad():
+                transcription_rnnt = asr_manager.model(wav, asr_manager.model_language[language], "rnnt")
+            return TranscriptionResponse(text=transcription_rnnt)
+        except Exception as e:
+            logger.error(f"Error in transcription: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
 
 @app.post("/v1/speech_to_speech")
 async def speech_to_speech(
@@ -867,28 +848,20 @@ async def speech_to_speech(
     file: UploadFile = File(...),
     language: str = Query(..., enum=list(asr_manager.model_language.keys())),
 ) -> StreamingResponse:
-    if not tts_manager.model:
-        raise HTTPException(status_code=503, detail="TTS model not loaded")
-    transcription = await transcribe_audio(file, language)
-    logger.info(f"Transcribed text: {transcription.text}")
+    async with request_queue:
+        if not tts_manager.model:
+            raise HTTPException(status_code=503, detail="TTS model not loaded")
+        transcription = await transcribe_audio(file, language)
+        logger.info(f"Transcribed text: {transcription.text}")
+        chat_request = ChatRequest(prompt=transcription.text, src_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"), tgt_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"))
+        processed_text = await chat(request, chat_request)
+        logger.info(f"Processed text: {processed_text.response}")
+        voice_request = KannadaSynthesizeRequest(text=processed_text.response)
+        audio_response = await synthesize_kannada(voice_request)
+        return audio_response
+
+LANGUAGE_TO_SCRIPT = {"kannada": "kan_Knda"}
 
-    chat_request = ChatRequest(
-        prompt=transcription.text,
-        src_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"),
-        tgt_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda")
-    )
-    processed_text = await chat(request, chat_request)
-    logger.info(f"Processed text: {processed_text.response}")
-
-    voice_request = KannadaSynthesizeRequest(text=processed_text.response)
-    audio_response = await synthesize_kannada(voice_request)
-    return audio_response
-
-LANGUAGE_TO_SCRIPT = {
-    "kannada": "kan_Knda"
-}
-
-# Main Execution
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run the FastAPI server.")
     parser.add_argument("--port", type=int, default=settings.port, help="Port to run the server on.")
@@ -915,15 +888,12 @@ if __name__ == "__main__":
     settings.speech_rate_limit = global_settings["speech_rate_limit"]
 
     llm_manager = LLMManager(settings.llm_model_name)
-
     if selected_config["components"]["ASR"]:
-        asr_model_name = selected_config["components"]["ASR"]["model"]
         asr_manager.model_language[selected_config["language"]] = selected_config["components"]["ASR"]["language_code"]
-
     if selected_config["components"]["Translation"]:
         translation_configs.extend(selected_config["components"]["Translation"])
 
     host = args.host if args.host != settings.host else settings.host
     port = args.port if args.port != settings.port else settings.port
 
-    uvicorn.run(app, host=host, port=port)
+    uvicorn.run(app, host=host, port=port, workers=2)