main.py

from contextlib import asynccontextmanager
from importlib.metadata import version
import asyncio
import base64
import io
import logging
import os
import sys
import tempfile
import time
import traceback

from fastapi import FastAPI, Request, Response, HTTPException
from fastapi.responses import JSONResponse

from typing import Dict, Any

# Configure logging with timestamps as a fallback
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)

def setup_uvicorn_logging():
    try:
        from uvicorn.logging import DefaultFormatter, AccessFormatter
        # Inject timestamps into Uvicorn's loggers so both server and app logs match
        for logger_name in ("uvicorn", "uvicorn.error", "uvicorn.access"):
            l = logging.getLogger(logger_name)
            for handler in l.handlers:
                use_colors = getattr(handler.formatter, "use_colors", True)
                if logger_name == "uvicorn.access":
                    handler.setFormatter(AccessFormatter(
                        fmt="%(asctime)s %(levelprefix)s %(client_addr)s - \"%(request_line)s\" %(status_code)s",
                        use_colors=use_colors,
                        datefmt="%Y-%m-%d %H:%M:%S"
                    ))
                else:
                    handler.setFormatter(DefaultFormatter(
                        fmt="%(asctime)s %(levelprefix)s %(message)s",
                        use_colors=use_colors,
                        datefmt="%Y-%m-%d %H:%M:%S"
                    ))
    except Exception as e:
        # Fallback silently if uvicorn is not running or logging structure is different
        pass

# Patch Uvicorn logging immediately on import
setup_uvicorn_logging()

# Use uvicorn's error logger if running under uvicorn, which handles formatting nicely
logger = logging.getLogger("uvicorn.error")

# this must come before attempts to import from transformers or torch
print(f'Versions: transformers: {version("transformers")}, torch: {version("torch")}')

import os
# Configure PyTorch to print clean logs when compilation starts and finishes
os.environ["TORCH_LOGS"] = "compiles"

import torch
# Enable TensorFloat32 (TF32) tensor cores for faster Float32 math on L4
torch.set_float32_matmul_precision('high')

# Increase the compiler's recompile limit to allow all transformer layers to compile.
# LLMs with KV Caches trigger one recompile per layer during the first token's
# initialization phase. A limit of 64 easily covers Qwen2's 28 layers.
import torch._dynamo
torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.recompile_limit = 64

import torchaudio
import torchaudio.transforms as T

# Bypass lazy-loader
from transformers.models.auto.processing_auto import AutoProcessor
from transformers.models.vibevoice_asr.modeling_vibevoice_asr import VibeVoiceAsrForConditionalGeneration

class EndpointHandler():
    def __init__(self, path=""):
        # 1. Load the processor
        self.processor = AutoProcessor.from_pretrained(path)

        # 2. Load and configure the model config to force Flash Attention 2 on the text decoder
        from transformers import AutoConfig
        config = AutoConfig.from_pretrained(path)
        if hasattr(config, "text_config"):
            config.text_config._attn_implementation = "flash_attention_2"
            logger.info("Forced Flash Attention 2 on the text decoder.")

        # 3. Load the specific VibeVoice model class in BF16
        self.model = VibeVoiceAsrForConditionalGeneration.from_pretrained(
            path,
            config=config,
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )

        # 4. Compile ONLY the Qwen2 text decoder to eliminate eager dequantization overhead
        logger.info("Compiling the Qwen2 text decoder with torch.compile(..., dynamic=True)...")
        compile_start = time.time()
        self.model.base_model.language_model = torch.compile(
            self.model.base_model.language_model,
            dynamic=True
        )
        logger.info(f"Model compilation wrapper set up in {time.time() - compile_start:.3f}s. Note: The very first request will trigger Triton compilation and take 2-3 minutes, but all subsequent requests will be blazing fast.")

        # Print layer device allocation
        if hasattr(self.model, "hf_device_map"):
            logger.info("Model layers device allocation (hf_device_map):")
            for layer_name, device in self.model.hf_device_map.items():
                logger.info(f"  - {layer_name}: {device}")
        else:
            # Fallback if hf_device_map is not populated
            devices = set()
            for name, param in self.model.named_parameters():
                devices.add(str(param.device))
            logger.info(f"Model loaded. Active devices for parameters: {list(devices)}")

        # 3. Dynamically fetch the expected sample rate (usually 16kHz or 24kHz)
        self.target_sr = getattr(self.processor.feature_extractor, "sampling_rate", 16000)

        # 4. Dynamically resolve all valid text and audio EOS (stopping) token IDs
        text_eos = self.processor.tokenizer.eos_token_id
        audio_eos = getattr(self.model.config, "audio_eos_token_id", 151647)
        if isinstance(text_eos, list):
            self.eos_token_ids = text_eos + [audio_eos]
        else:
            self.eos_token_ids = [text_eos, audio_eos]
        logger.info(f"Configured stopping tokens (eos_token_id): {self.eos_token_ids}")

        # 5. Warm up the compiled model to trigger Triton compilation at startup
        logger.info("Warming up the compiled model (this will trigger Triton compilation and take ~1.5 - 2 minutes)...")
        warmup_start = time.time()
        try:
            import numpy as np
            from transformers.cache_utils import StaticCache
            
            # 1 second of silence
            dummy_audio = np.zeros(self.target_sr)
            processed_inputs = self.processor(
                text="transcribe", # Must be non-empty so sequence length > 0 (prevents compiler shape ambiguity)
                audio=dummy_audio,
                prompt="transcribe"
            )
            # Move to device and dtype
            for k, v in processed_inputs.items():
                if isinstance(v, torch.Tensor):
                    if torch.is_floating_point(v):
                        processed_inputs[k] = v.to(device=self.model.device, dtype=self.model.dtype)
                    else:
                        processed_inputs[k] = v.to(device=self.model.device)
            
            # Manually allocate the 16k StaticCache to pre-compile the 1-hour production shape
            cache = StaticCache(
                config=self.model.base_model.language_model.config,
                max_batch_size=1,
                max_cache_len=16384, # Match the 1-hour production shape!
                device=self.model.device,
                dtype=self.model.dtype
            )
            
            # Run a dummy generate pass, but stop after 5 tokens to keep boot fast!
            with torch.no_grad():
                _ = self.model.generate(
                    **processed_inputs,
                    max_new_tokens=5,
                    past_key_values=cache,
                    eos_token_id=self.eos_token_ids,
                    repetition_penalty=1.1,
                    no_repeat_ngram_size=5
                )
            logger.info(f"Model warmed up and Triton kernels compiled successfully in {time.time() - warmup_start:.3f}s! Server is now ready to handle requests instantly.")
        except Exception as e:
            logger.warning(f"Failed to warm up/compile model during startup: {e}. Compilation will happen on the first request instead.")

    def __call__(self, data: Any) -> Dict[str, Any]:
        start_time = time.time()
        # 1. Extract payload (Handles both JSON payloads and raw binary uploads)
        if isinstance(data, dict):
            data_copy = data.copy()
            inputs = data_copy.pop("inputs", None)
            if inputs is None:
                inputs = data_copy
            parameters = data_copy.pop("parameters", {})
        else:
            inputs = data
            parameters = {}

        if not inputs:
            return {"error": "Missing 'inputs' in request data"}

        hotwords = parameters.get("hotwords", None)
        return_format = parameters.get("return_format", "parsed") # Default to parsed for rich output

        temp_file = None
        audio_array = None

        # 2. Decode raw audio bytes (various formats) using torchaudio
        logger.info("Starting audio loading and preprocessing...")
        audio_load_start = time.time()

        def load_and_resample(audio_path):
            waveform, sample_rate = torchaudio.load(audio_path)
            if waveform.shape[0] > 1:
                # Mixdown stereo to mono
                waveform = waveform.mean(dim=0, keepdim=True)
            if sample_rate != self.target_sr:
                resampler = T.Resample(orig_freq=sample_rate, new_freq=self.target_sr)
                waveform = resampler(waveform)
            return waveform.squeeze().numpy()

        try:
            if isinstance(inputs, bytes):
                logger.info(f"Decoding audio from raw bytes ({len(inputs)} bytes)...")
                # Raw bytes from binary upload (could be MP3, WAV, FLAC, etc.)
                # Write to suffix-less temp file so torchaudio can load it
                temp_file = tempfile.NamedTemporaryFile("wb", delete=False)
                temp_file.write(inputs)
                temp_file.flush()
                temp_file.close()

                audio_array = load_and_resample(temp_file.name)

            elif isinstance(inputs, str):
                if inputs.startswith("http://") or inputs.startswith("https://"):
                    logger.info(f"Downloading audio from URL: {inputs}...")
                    # URL input - download first to be safe
                    import requests
                    response = requests.get(inputs)
                    temp_file = tempfile.NamedTemporaryFile("wb", delete=False)
                    temp_file.write(response.content)
                    temp_file.flush()
                    temp_file.close()

                    audio_array = load_and_resample(temp_file.name)
                else:
                    logger.info("Decoding audio from Base64 string...")
                    # Try base64 decode
                    try:
                        decoded_bytes = base64.b64decode(inputs)
                        temp_file = tempfile.NamedTemporaryFile("wb", delete=False)
                        temp_file.write(decoded_bytes)
                        temp_file.flush()
                        temp_file.close()

                        audio_array = load_and_resample(temp_file.name)
                    except Exception as e:
                        logger.info(f"Base64 decode failed ({e}), assuming input is a local file path...")
                        # Fallback to assuming it's a local path
                        audio_array = load_and_resample(inputs)
            else:
                logger.info("Using pre-loaded audio array...")
                # If already loaded (e.g. numpy array passed in some test environments)
                audio_array = inputs

            audio_load_duration = time.time() - audio_load_start
            audio_duration_sec = len(audio_array) / self.target_sr if audio_array is not None else 0
            logger.info(f"Audio loaded successfully in {audio_load_duration:.3f}s. "
                        f"Audio duration: {audio_duration_sec:.2f}s, Sample Rate: {self.target_sr}Hz")

            # 3. Prepare inputs using the recommended API
            logger.info("Preprocessing audio features...")
            preprocess_start = time.time()
            processed_inputs = self.processor.apply_transcription_request(
                audio=audio_array,
                prompt=hotwords
            )

            # Safely move to device and cast ONLY floating point tensors to the model's dtype.
            # Casting integer tensors (like input_ids) to bfloat16/float16 will cause model errors.
            for k, v in processed_inputs.items():
                if isinstance(v, torch.Tensor):
                    if torch.is_floating_point(v):
                        processed_inputs[k] = v.to(device=self.model.device, dtype=self.model.dtype)
                    else:
                        processed_inputs[k] = v.to(device=self.model.device)
            logger.info(f"Preprocessing completed in {time.time() - preprocess_start:.3f}s.")

            # 4. Generate
            logger.info("Starting model inference (generation)...")
            inference_start = time.time()
            with torch.no_grad():
                output_ids = self.model.generate(
                    **processed_inputs,
                    # CRITICAL: VibeVoice needs a token limit to handle up to 1-hour audio.
                    max_new_tokens=16384,
                    cache_implementation="static",
                    eos_token_id=self.eos_token_ids,
                    repetition_penalty=1.1,
                    no_repeat_ngram_size=5
                )
            inference_duration = time.time() - inference_start
            
            # Calculate token metrics for diagnostics
            num_input_tokens = processed_inputs["input_ids"].shape[1] if "input_ids" in processed_inputs else 0
            num_total_tokens = output_ids.shape[1]
            num_generated_tokens = num_total_tokens - num_input_tokens
            tokens_per_sec = num_generated_tokens / inference_duration if inference_duration > 0 else 0
            
            logger.info(f"Model inference completed in {inference_duration:.3f}s.")
            logger.info(f"Generated {num_generated_tokens} tokens (Input: {num_input_tokens}, Total: {num_total_tokens}).")
            logger.info(f"Generation speed: {tokens_per_sec:.2f} tokens/second.")

            # Slice generated IDs to exclude the prompt (Fixes prompt leakage)
            if "input_ids" in processed_inputs:
                prompt_len = processed_inputs["input_ids"].shape[1]
                generated_ids = output_ids[:, prompt_len:]
            else:
                generated_ids = output_ids

            # 5. Decode using return_format
            logger.info(f"Decoding generated tokens to text (format: '{return_format}')...")
            decode_start = time.time()
            try:
                transcription = self.processor.decode(generated_ids, return_format=return_format)[0]
            except Exception as e:
                logger.warning(f"Decoding with return_format='{return_format}' failed, falling back to batch_decode. Error: {e}")
                # Fallback to standard decode if return_format fails
                transcription = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
            logger.info(f"Token decoding completed in {time.time() - decode_start:.3f}s.")

            total_duration = time.time() - start_time
            logger.info(f"Request processed successfully in {total_duration:.3f}s.")
            return {"result": transcription}

        except Exception as e:
            logger.exception("Inference failed due to exception", e)
            return {"error": f"Inference failed: {str(e)}"}

        finally:
            # Clean up temp file
            if temp_file and os.path.exists(temp_file.name):
                try:
                    os.unlink(temp_file.name)
                except OSError:
                    pass

MODEL_DIR = os.environ.get("MODEL_DIR", "/repository")
handler = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global handler
    logger.info(f"Loading model from {MODEL_DIR}...")
    handler = EndpointHandler(path=MODEL_DIR)
    logger.info("Model loaded successfully.")
    yield
    del handler

app = FastAPI(lifespan=lifespan)

@app.get("/")
@app.get("/health")
def health_check():
    return {"status": "ok"}

@app.post("/")
async def predict(request: Request):
    if handler is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    content_type = request.headers.get("content-type", "")

    if "application/json" in content_type:
        try:
            data = await request.json()
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}")
    else:
        # Fallback to raw bytes
        data = await request.body()

    response = await asyncio.to_thread(handler, data)

    if "error" in response:
        raise HTTPException(status_code=500, detail=response["error"])

    return response

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    # Force print to stderr
    print(f"CRITICAL ERROR: {str(exc)}", file=sys.stderr)
    print(traceback.format_exc(), file=sys.stderr)

    # Return the error in the payload so you can see it in your HTTP client
    return JSONResponse(
        status_code=500,
        content={"error": str(exc), "type": str(type(exc))}
    )