# SPDX-License-Identifier: Apache-2.0
# Derived from vLLM; modified by DeepSeek.

from __future__ import annotations

import importlib

import torch.nn as nn

from deepinfer.logger import init_logger

logger = init_logger(__name__)

# Architecture -> (module, class).
_MODELS = {
    "DeepseekV41ForCausalLM": ("deepseek_vl", "DeepseekVLForCausalLM"),
}


class ModelRegistry:
    @staticmethod
    def load_model_cls(model_arch: str) -> type[nn.Module] | None:
        if model_arch not in _MODELS:
            return None

        module_name, model_cls_name = _MODELS[model_arch]
        module = importlib.import_module(f"deepinfer.model_executor.models.{module_name}")
        return getattr(module, model_cls_name, None)

    @staticmethod
    def get_supported_archs() -> list[str]:
        return list(_MODELS.keys())


__all__ = [
    "ModelRegistry",
]