Class AutoModelVLLM

funasr/auto/auto_model_vllm.py:181–309

Generic vLLM wrapper for LLM-based ASR models. Automatically detects model architecture, extracts LLM weights, loads audio components in PyTorch, and uses vLLM for generation. Works for: FunASRNano, LLMASR, GLMASR, and any model with audio_encoder + audio_adaptor + LLM architecture

class AutoModelVLLM:
    """Generic vLLM wrapper for LLM-based ASR models.

    Automatically detects model architecture, extracts LLM weights,
    loads audio components in PyTorch, and uses vLLM for generation.

    Works for: FunASRNano, LLMASR, GLMASR, and any model with
    audio_encoder + audio_adaptor + LLM architecture.

    Args:
        model: Model name (hub) or local directory path.
        hub: "ms" (ModelScope) or "hf" (HuggingFace).
        device: Device for audio encoder/adaptor.
        dtype: Compute dtype ("bf16", "fp16", "fp32").
        tensor_parallel_size: GPUs for vLLM tensor parallelism.
        gpu_memory_utilization: GPU memory fraction for vLLM.
        max_model_len: Maximum sequence length.

    Example:
        >>> model = AutoModelVLLM(model="FunAudioLLM/Fun-ASR-Nano-2512")
        >>> results = model.generate(["audio.wav"], language="中文")
    """

    def __init__(
        self,
        model: str,
        hub: str = "ms",
        device: str = "cuda:0",
        dtype: str = "bf16",
        tensor_parallel_size: int = 1,
        gpu_memory_utilization: float = 0.8,
        max_model_len: int = 4096,
        enforce_eager: bool = False,
        **kwargs,
    ):
        # Resolve model directory
        if os.path.isdir(model):
            self.model_dir = model
        else:
            if hub in ("ms", "modelscope"):
                from modelscope.hub.snapshot_download import snapshot_download
                self.model_dir = snapshot_download(model, revision=kwargs.get("revision", "master"))
            elif hub in ("hf", "huggingface"):
                from huggingface_hub import snapshot_download
                self.model_dir = snapshot_download(model)
            else:
                raise ValueError(f"Unsupported hub: {hub}")

        # Check model type
        from omegaconf import OmegaConf
        config = OmegaConf.load(os.path.join(self.model_dir, "config.yaml"))
        self.model_type = config.get("model", "unknown")
        check_vllm_applicable(self.model_type)

        self.device = device
        self.dtype = dtype
        self.torch_dtype = dtype_map.get(dtype, torch.bfloat16)