import whisperx

_asr_model = None
_alignment_model = None
_metadata = None

def load_asr(device="cpu", model_size="medium"):
    global _asr_model
    if _asr_model is None:
        _asr_model = whisperx.load_model(model_size, device=device, compute_type="float32")
    return _asr_model

def transcribe_and_align(audio_path: str, device="cpu"):
    model = load_asr(device=device)
    result = model.transcribe(audio_path, batch_size=16)

    global _alignment_model, _metadata
    if _alignment_model is None or _metadata is None:
        _alignment_model, _metadata = whisperx.load_align_model(
            language_code=result["language"], device=device
        )
    # Pass device here as required
    result_aligned = whisperx.align(
        result["segments"], _alignment_model, _metadata, audio_path, device=device
    )
    return result, result_aligned