# SPDX-License-Identifier: Apache-2.0
"""
Central registry for multimodal models.

This module provides a centralized registry for multimodal models, including pipelines
and sampling parameters. It allows for easy registration and retrieval of model
information based on model paths or other identifiers.
"""

import dataclasses
import importlib
import os
import pkgutil
from functools import lru_cache
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Tuple,
    Type,
    Union,
)

if TYPE_CHECKING:
    from sglang.multimodal_gen.runtime.server_args import Backend

from sglang.multimodal_gen.configs.pipeline_configs import (
    Cosmos3Config,
    FastHunyuanConfig,
    FluxPipelineConfig,
    HeliosDistilledConfig,
    HeliosMidConfig,
    HeliosT2VConfig,
    HunyuanConfig,
    LingBotWorldCausalDMDConfig,
    LingBotWorldV2CausalDMDConfig,
    WanI2V480PConfig,
    WanI2V720PConfig,
    WanT2V480PConfig,
    WanT2V720PConfig,
    ZImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.ernie_image import (
    ErnieImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.flux import (
    Flux2KleinBasePipelineConfig,
    Flux2KleinPipelineConfig,
    Flux2PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.glm_image import (
    GlmImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
    Hunyuan3D2PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.ideogram import (
    Ideogram4PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import (
    JoyEchoPipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.joy_image import (
    JoyImageEditPipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.krea2 import Krea2PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
    LTX2PipelineConfig,
    LTX23PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.mova import (
    MOVA360PConfig,
    MOVA720PConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.pi05 import Pi05PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
    QwenImageEditPipelineConfig,
    QwenImageEditPlus_2511_PipelineConfig,
    QwenImageEditPlusPipelineConfig,
    QwenImageLayeredPipelineConfig,
    QwenImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import (
    StableDiffusion3PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.wan import (
    FastWan2_1_T2V_480P_Config,
    FastWan2_2_TI2V_5B_Config,
    TurboWanI2V720Config,
    TurboWanT2V1_3B480PConfig,
    TurboWanT2V480PConfig,
    Wan2_2_I2V_A14B_Config,
    Wan2_2_T2V_A14B_Config,
    Wan2_2_TI2V_5B_Config,
)
from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams
from sglang.multimodal_gen.configs.sample.ernie_image import ErnieImageSamplingParams
from sglang.multimodal_gen.configs.sample.flux import (
    Flux2KleinBaseSamplingParams,
    Flux2KleinSamplingParams,
    Flux2SamplingParams,
    FluxSamplingParams,
)
from sglang.multimodal_gen.configs.sample.glmimage import GlmImageSamplingParams
from sglang.multimodal_gen.configs.sample.helios import (
    HeliosDistilledSamplingParams,
    HeliosMidSamplingParams,
    HeliosT2VSamplingParams,
)
from sglang.multimodal_gen.configs.sample.hunyuan import (
    FastHunyuanSamplingParam,
    HunyuanSamplingParams,
)
from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams
from sglang.multimodal_gen.configs.sample.ideogram import Ideogram4SamplingParams
from sglang.multimodal_gen.configs.sample.joy_echo import JoyEchoSamplingParams
from sglang.multimodal_gen.configs.sample.joy_image import (
    JoyImageEditSamplingParams,
)
from sglang.multimodal_gen.configs.sample.krea2 import (
    Krea2SamplingParams,
)
from sglang.multimodal_gen.configs.sample.lingbot_world import (
    LingBotWorldSamplingParams,
)
from sglang.multimodal_gen.configs.sample.ltx_2 import (
    LTX2SamplingParams,
    LTX23HQSamplingParams,
    LTX23SamplingParams,
)
from sglang.multimodal_gen.configs.sample.mova import (
    MOVA_360P_SamplingParams,
    MOVA_720P_SamplingParams,
)
from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
from sglang.multimodal_gen.configs.sample.qwenimage import (
    QwenImage2512SamplingParams,
    QwenImageEditPlusSamplingParams,
    QwenImageLayeredSamplingParams,
    QwenImageSamplingParams,
)
from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams
from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams
from sglang.multimodal_gen.configs.sample.stablediffusion3 import (
    StableDiffusion3SamplingParams,
)
from sglang.multimodal_gen.configs.sample.wan import (
    FastWanT2V480PConfig,
    Turbo_Wan2_2_I2V_A14B_SamplingParam,
    Wan2_1_Fun_1_3B_InP_SamplingParams,
    Wan2_2_I2V_A14B_SamplingParam,
    Wan2_2_T2V_A14B_SamplingParam,
    Wan2_2_TI2V_5B_SamplingParam,
    WanI2V_14B_480P_SamplingParam,
    WanI2V_14B_720P_SamplingParam,
    WanT2V_1_3B_SamplingParams,
    WanT2V_14B_SamplingParams,
)
from sglang.multimodal_gen.configs.sample.zimage import (
    ZImageSamplingParams,
    ZImageTurboSamplingParams,
)
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
    ComposedPipelineBase,
)
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
    maybe_download_model_index,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.utils import KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS

logger = init_logger(__name__)

# --- Part 1: Pipeline Discovery ---

_PIPELINE_REGISTRY: Dict[str, Type[ComposedPipelineBase]] = {}

# Registry for pipeline configuration classes (for safetensors files without model_index.json)
# Maps pipeline_class_name -> (PipelineConfig class, SamplingParams class)
_PIPELINE_CONFIG_REGISTRY: Dict[str, Tuple[Type[PipelineConfig], Type[Any]]] = {}


def _discover_and_register_pipelines():
    """
    Automatically discover and register all ComposedPipelineBase subclasses.
    This function scans the 'sglang.multimodal_gen.runtime.pipelines' package,
    finds modules with an 'EntryClass' attribute, and maps the class's 'pipeline_name'
    to the class itself in a global registry.
    """
    if _PIPELINE_REGISTRY:  # run only once
        return

    package_name = "sglang.multimodal_gen.runtime.pipelines"
    package = importlib.import_module(package_name)

    for _, module_name, ispkg in pkgutil.walk_packages(
        package.__path__, package.__name__ + "."
    ):
        if not ispkg:
            try:
                pipeline_module = importlib.import_module(module_name)
            except Exception as exc:
                logger.warning(
                    "Skipping pipeline module %s during discovery due to import failure: %s",
                    module_name,
                    exc,
                )
                logger.debug(
                    "Pipeline import failure details for %s", module_name, exc_info=True
                )
                continue
            if hasattr(pipeline_module, "EntryClass"):
                entry_cls = pipeline_module.EntryClass
                entry_cls_list = (
                    [entry_cls] if not isinstance(entry_cls, list) else entry_cls
                )

                for cls in entry_cls_list:
                    if not issubclass(cls, ComposedPipelineBase):
                        continue
                    if cls.pipeline_name in _PIPELINE_REGISTRY:
                        logger.warning(
                            f"Duplicate pipeline name '{cls.pipeline_name}' found. Overwriting."
                        )
                    _PIPELINE_REGISTRY[cls.pipeline_name] = cls

                    # Special handling for ComfyUI Pipelines:
                    # Auto-register config classes if Pipeline class has them defined
                    # since comfyui get model from a single weight file, so we need to register the config classes here
                    if hasattr(cls, "pipeline_config_cls") and hasattr(
                        cls, "sampling_params_cls"
                    ):
                        _PIPELINE_CONFIG_REGISTRY[cls.pipeline_name] = (
                            cls.pipeline_config_cls,
                            cls.sampling_params_cls,
                        )
                        logger.debug(
                            f"Auto-registered config classes for pipeline '{cls.pipeline_name}': "
                            f"PipelineConfig={cls.pipeline_config_cls.__name__}, "
                            f"SamplingParams={cls.sampling_params_cls.__name__}"
                        )
    logger.debug(
        f"Registering pipelines complete, {len(_PIPELINE_REGISTRY)} pipelines registered"
    )


def get_pipeline_config_classes(
    pipeline_class_name: str,
) -> Tuple[Type[PipelineConfig], Type[Any]] | None:
    """
    Get the configuration classes for a pipeline.
    """
    # Ensure pipelines are discovered first
    _discover_and_register_pipelines()
    return _PIPELINE_CONFIG_REGISTRY.get(pipeline_class_name)


# --- Part 2: Config Registration ---
@dataclasses.dataclass
class ConfigInfo:
    """Encapsulates all configuration information required to register a
    diffusers model within this framework."""

    sampling_param_cls: Any
    pipeline_config_cls: Type[PipelineConfig]


# The central registry mapping a model name to its configuration information
_CONFIG_REGISTRY: Dict[str, ConfigInfo] = {}

# Mappings from Hugging Face model paths to our internal model names
_MODEL_HF_PATH_TO_NAME: Dict[str, str] = {}

# Detectors to identify model families from paths or class names
_MODEL_NAME_DETECTORS: List[Tuple[str, Callable[[str], bool]]] = []


def register_configs(
    sampling_param_cls: Any,
    pipeline_config_cls: Type[PipelineConfig],
    hf_model_paths: Optional[List[str]] = None,
    model_detectors: Optional[List[Callable[[str], bool]]] = None,
):
    """
    Registers configuration classes for a new model family.
    """
    model_id = str(len(_CONFIG_REGISTRY))

    _CONFIG_REGISTRY[model_id] = ConfigInfo(
        sampling_param_cls=sampling_param_cls,
        pipeline_config_cls=pipeline_config_cls,
    )
    if hf_model_paths:
        for path in hf_model_paths:
            if path in _MODEL_HF_PATH_TO_NAME:
                logger.warning(
                    f"Model path '{path}' is already mapped to '{_MODEL_HF_PATH_TO_NAME[path]}' and will be overwritten by '{model_id}'."
                )
            _MODEL_HF_PATH_TO_NAME[path] = model_id

    if model_detectors:
        for detector in model_detectors:
            _MODEL_NAME_DETECTORS.append((model_id, detector))


def get_model_short_name(model_id: str) -> str:
    if "/" in model_id:
        return model_id.rstrip("/").split("/")[-1]
    else:
        return model_id


def _normalize_hf_cache_path(path: str) -> str:
    """Normalize a local HuggingFace cache path before substring matching.

    We match registered repo ids like ``org/repo`` against cache fragments like ``models--org--repo`` that appear in snapshot/blob paths.
    """
    return os.path.normpath(path).lower().replace("\\", "/")


def has_registered_diffusion_model_path(model_path: str) -> bool:
    all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True)

    if model_path in _MODEL_HF_PATH_TO_NAME:
        return True

    model_short_name = get_model_short_name(model_path.lower())
    for registered_model_hf_id in all_model_hf_paths:
        registered_model_name = get_model_short_name(registered_model_hf_id.lower())
        if registered_model_name in model_short_name:
            return True

    normalized_model_path = _normalize_hf_cache_path(model_path)
    for registered_model_hf_id in all_model_hf_paths:
        cache_repo_fragment = (
            f"models--{registered_model_hf_id.lower().replace('/', '--')}"
        )
        if cache_repo_fragment in normalized_model_path:
            return True

    return False


@lru_cache(maxsize=1)
def _get_config_info(
    model_path: str, model_id: Optional[str] = None
) -> Optional[ConfigInfo]:
    """
    Gets the ConfigInfo for a given model path using mappings and detectors.
    """
    all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True)

    # 0. Explicit model_id override: match by short name
    if model_id is not None:
        model_id_lower = model_id.lower()
        for registered_hf_id in all_model_hf_paths:
            if get_model_short_name(registered_hf_id).lower() == model_id_lower:
                logger.debug(
                    f"Resolved model via explicit --model-id '{model_id}' → '{registered_hf_id}'."
                )
                return _CONFIG_REGISTRY.get(_MODEL_HF_PATH_TO_NAME[registered_hf_id])
        logger.warning(
            f"--model-id '{model_id}' did not match any registered model; "
            "falling back to automatic detection."
        )

    # 1. Exact match
    if model_path in _MODEL_HF_PATH_TO_NAME:
        model_id = _MODEL_HF_PATH_TO_NAME[model_path]
        logger.debug(f"Resolved model path '{model_path}' from exact path match.")
        return _CONFIG_REGISTRY.get(model_id)

    # 2. Partial match: find the best (longest) match against all registered model hf paths.
    model_short_name = get_model_short_name(model_path.lower())
    for registered_model_hf_id in all_model_hf_paths:
        registered_model_name = get_model_short_name(registered_model_hf_id.lower())

        if registered_model_name in model_short_name:
            logger.debug(
                f"Resolved model name '{registered_model_hf_id}' from partial path match."
            )
            model_id = _MODEL_HF_PATH_TO_NAME[registered_model_hf_id]
            return _CONFIG_REGISTRY.get(model_id)

    # 2b. Match local HuggingFace cache snapshot/blob paths such as:
    #   .../models--org--repo/snapshots/<hash>
    # This lets users pass a local HF cache snapshot directory directly even
    # when its basename is only the snapshot hash.
    # Example:
    #    /xxx/models--black-forest-labs--FLUX.2-dev-NVFP4/snapshots/142b87e70bc3006937b7093d89ff287b5f59f071
    # -> models--black-forest-labs--flux.2-dev-nvfp4 (to match with cache_repo_fragment)
    normalized_model_path = _normalize_hf_cache_path(model_path)
    for registered_model_hf_id in all_model_hf_paths:
        cache_repo_fragment = (
            f"models--{registered_model_hf_id.lower().replace('/', '--')}"
        )
        if cache_repo_fragment in normalized_model_path:
            logger.debug(
                "Resolved HuggingFace cache path '%s' to registered model '%s'.",
                model_path,
                registered_model_hf_id,
            )
            model_id = _MODEL_HF_PATH_TO_NAME[registered_model_hf_id]
            return _CONFIG_REGISTRY.get(model_id)

    # 3. Use detectors
    config = maybe_download_model_index(model_path)
    pipeline_name = config.get("_class_name", "").lower()

    matched_model_names = []
    for model_id, detector in _MODEL_NAME_DETECTORS:
        if detector(model_path.lower()) or detector(pipeline_name):
            logger.debug(
                f"Matched model name '{model_id}' using a registered detector."
            )
            matched_model_names += [model_id]

    if len(matched_model_names) >= 1:
        if len(matched_model_names) > 1:
            logger.warning(
                "More than one model name is matched, using the first matched"
            )
        model_id = matched_model_names[0]
        return _CONFIG_REGISTRY.get(model_id)
    else:
        logger.debug(
            f"No model info found for model path: {model_path}. "
            f"Please check the model path or specify the model_id explicitly."
        )
        return None


# --- Part 3: Main Resolver ---


@dataclasses.dataclass
class ModelInfo:
    """
    Encapsulates all configuration information required to register a
    diffusers model within this framework.
    """

    pipeline_cls: Type[ComposedPipelineBase]
    sampling_param_cls: Any
    pipeline_config_cls: Type[PipelineConfig]


def _get_diffusers_model_info(
    model_path: Optional[str] = None,
    model_id: Optional[str] = None,
) -> ModelInfo:
    """
    Get model info for diffusers backend.

    Returns a ModelInfo with DiffusersPipeline and generic configs.
    When model_path is provided and has a registered native config,
    inherits task_type from it so that validation (e.g. accepts_image_input)
    works correctly even under the diffusers backend.
    """
    from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import (
        DIFFUSERS_TASK_TYPE_TO_CONFIG,
        DiffusersGenericPipelineConfig,
    )
    from sglang.multimodal_gen.configs.sample.diffusers_generic import (
        DiffusersGenericSamplingParams,
    )
    from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import (
        DiffusersPipeline,
    )

    sampling_param_cls = DiffusersGenericSamplingParams
    pipeline_config_cls = DiffusersGenericPipelineConfig

    # If there is a registered native config for this model, inherit its task_type.
    # We use pre-defined static subclasses instead of make_dataclass so the config
    # class is pickle-safe for multiprocessing spawn (fixes #21453).
    if model_path is not None:
        config_info = _get_config_info(model_path, model_id=model_id)
        if config_info is not None:
            sampling_param_cls = config_info.sampling_param_cls
            native_task_type = config_info.pipeline_config_cls.task_type
            if native_task_type != DiffusersGenericPipelineConfig.task_type:
                pipeline_config_cls = DIFFUSERS_TASK_TYPE_TO_CONFIG.get(
                    native_task_type, DiffusersGenericPipelineConfig
                )
                logger.debug(
                    "Inherited task_type=%s from native config for diffusers backend",
                    native_task_type.name,
                )

    return ModelInfo(
        pipeline_cls=DiffusersPipeline,
        sampling_param_cls=sampling_param_cls,
        pipeline_config_cls=pipeline_config_cls,
    )


@lru_cache(maxsize=1)
def get_model_info(
    model_path: str,
    backend: Optional[Union[str, "Backend"]] = None,
    model_id: Optional[str] = None,
) -> Optional[ModelInfo]:
    """
    Resolves all necessary classes (pipeline, sampling, config) for a given model path.

    This function serves as the main entry point for model resolution. It performs two main tasks:
    1. Dynamically resolves the pipeline class by reading 'model_index.json' and matching
       '_class_name' against an auto-discovered registry of pipeline implementations.
    2. Resolves the associated configuration classes (for sampling and pipeline) using a
       manually registered mapping based on the model path.

    Args:
        backend: Backend to use ('auto', 'sglang', 'diffusers'). If None, uses 'auto'.

    """
    # import Backend enum here to avoid circular imports
    from sglang.multimodal_gen.runtime.server_args import Backend

    # Normalize backend
    if backend is None:
        backend = Backend.AUTO
    elif isinstance(backend, str):
        backend = Backend.from_string(backend)

    # Handle explicit diffusers backend
    if backend == Backend.DIFFUSERS:
        logger.info(
            "Using diffusers backend for model '%s' (explicitly requested)", model_path
        )
        return _get_diffusers_model_info(model_path=model_path, model_id=model_id)

    # For AUTO or SGLANG backend, try native implementation first
    # 1. Discover all available pipeline classes and cache them
    _discover_and_register_pipelines()

    # Detect quantized models and fallback to diffusers
    is_quantized = any(q in model_path.lower() for q in ["-4bit", "-awq", "-gptq"])
    if is_quantized and backend != Backend.DIFFUSERS:
        logger.info(
            "Detected a quantized model format ('%s'). "
            "The native sglang-diffusion engine currently only supports BF16/FP16. "
            "Falling back to diffusers backend.",
            model_path,
        )
        return _get_diffusers_model_info(model_path=model_path, model_id=model_id)

    # 2. Get pipeline class - check non-diffusers models first
    pipeline_class_name = get_non_diffusers_pipeline_name(model_path)
    if pipeline_class_name:
        # Known non-diffusers model, skip model_index.json download
        logger.debug(
            f"Using registered pipeline '{pipeline_class_name}' for non-diffusers model '{model_path}'"
        )
    else:
        # Try to get from model_index.json
        try:
            config = maybe_download_model_index(model_path)
        except Exception as e:
            logger.error(f"Could not read model config for '{model_path}': {e}")
            if backend == Backend.AUTO:
                logger.info("Falling back to diffusers backend")
                return _get_diffusers_model_info(
                    model_path=model_path, model_id=model_id
                )
            return None

        pipeline_class_name = config.get("_class_name")
        if not pipeline_class_name:
            logger.error(
                f"'_class_name' not found in model_index.json for '{model_path}'"
            )
            if backend == Backend.AUTO:
                logger.info("Falling back to diffusers backend")
                return _get_diffusers_model_info(
                    model_path=model_path, model_id=model_id
                )
            return None

    pipeline_cls = _PIPELINE_REGISTRY.get(pipeline_class_name)
    if not pipeline_cls:
        if backend == Backend.AUTO:
            logger.warning(
                f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' has no native sglang support. "
                f"Falling back to diffusers backend."
            )
            return _get_diffusers_model_info(model_path=model_path, model_id=model_id)
        else:
            logger.error(
                f"Pipeline class '{pipeline_class_name}' specified in '{model_path}' is not a registered EntryClass in the framework. "
                f"Available pipelines: {list(_PIPELINE_REGISTRY.keys())}. "
                f"Consider using --backend diffusers to use vanilla diffusers pipeline."
            )
            return None

    # 3. Get configuration classes (sampling, pipeline config)
    config_info = _get_config_info(model_path, model_id=model_id)
    if not config_info:
        if backend == Backend.AUTO:
            logger.warning(
                f"Could not resolve native configuration for model '{model_path}'. "
                f"Falling back to diffusers backend."
            )
            return _get_diffusers_model_info(model_path=model_path, model_id=model_id)
        else:
            logger.error(
                f"Could not resolve configuration for model '{model_path}'. "
                "It is not a registered model path or detected by any registered model family detectors. "
                f"Known model paths: {list(_MODEL_HF_PATH_TO_NAME.keys())}. "
                f"Consider using --backend diffusers to use vanilla diffusers pipeline."
            )
            return None

    # 4. Combine and return the complete model info
    logger.debug("Using native sglang backend for model '%s'", model_path)
    model_info = ModelInfo(
        pipeline_cls=pipeline_cls,
        sampling_param_cls=config_info.sampling_param_cls,
        pipeline_config_cls=config_info.pipeline_config_cls,
    )
    logger.debug(f"Found model info: {model_info}")

    return model_info


# Registration of model configs
def _register_configs():
    # Pi0.5 / OpenPI / LeRobot action policies.
    register_configs(
        sampling_param_cls=Pi05SamplingParams,
        pipeline_config_cls=Pi05PipelineConfig,
        hf_model_paths=[
            "lerobot/pi05_base",
            "lerobot/pi05_libero_base",
        ],
        model_detectors=[
            lambda hf_id: "pi05" in hf_id.lower(),
            lambda hf_id: "pi0.5" in hf_id.lower(),
        ],
    )

    # LTX-2
    register_configs(
        sampling_param_cls=LTX2SamplingParams,
        pipeline_config_cls=LTX2PipelineConfig,
        hf_model_paths=["Lightricks/LTX-2"],
        model_detectors=[
            lambda path: "ltx" in path.lower() and "video" in path.lower(),
            lambda path: "ltx-2" in path.lower() and "ltx-2.3" not in path.lower(),
        ],
    )
    register_configs(
        sampling_param_cls=LTX23SamplingParams,
        pipeline_config_cls=LTX23PipelineConfig,
        hf_model_paths=["Lightricks/LTX-2.3"],
        model_detectors=[
            lambda path: "ltx-2.3" in path.lower(),
        ],
    )
    # register dedicated sampling params for LTX2TwoStageHQPipeline
    _PIPELINE_CONFIG_REGISTRY.setdefault(
        "LTX2TwoStageHQPipeline",
        (LTX2PipelineConfig, LTX23HQSamplingParams),
    )

    # Hunyuan
    register_configs(
        sampling_param_cls=HunyuanSamplingParams,
        pipeline_config_cls=HunyuanConfig,
        hf_model_paths=[
            "hunyuanvideo-community/HunyuanVideo",
        ],
        model_detectors=[lambda hf_id: "hunyuanvideo" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=FastHunyuanSamplingParam,
        pipeline_config_cls=FastHunyuanConfig,
        hf_model_paths=[
            "FastVideo/FastHunyuan-diffusers",
        ],
    )
    # Wan
    register_configs(
        sampling_param_cls=WanT2V_1_3B_SamplingParams,
        pipeline_config_cls=WanT2V480PConfig,
        hf_model_paths=[
            "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        ],
        model_detectors=[lambda hf_id: "wanpipeline" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=WanT2V_1_3B_SamplingParams,
        pipeline_config_cls=TurboWanT2V1_3B480PConfig,
        hf_model_paths=[
            "IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=WanT2V_14B_SamplingParams,
        pipeline_config_cls=WanT2V720PConfig,
        hf_model_paths=[
            "Wan-AI/Wan2.1-T2V-14B-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=WanT2V_14B_SamplingParams,
        pipeline_config_cls=TurboWanT2V480PConfig,
        hf_model_paths=[
            "IPostYellow/TurboWan2.1-T2V-14B-Diffusers",
            "IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=WanI2V_14B_480P_SamplingParam,
        pipeline_config_cls=WanI2V480PConfig,
        hf_model_paths=[
            "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
        ],
        model_detectors=[lambda hf_id: "wanimagetovideo" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=WanI2V_14B_720P_SamplingParam,
        pipeline_config_cls=WanI2V720PConfig,
        hf_model_paths=[
            "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=Turbo_Wan2_2_I2V_A14B_SamplingParam,
        pipeline_config_cls=TurboWanI2V720Config,
        hf_model_paths=[
            "IPostYellow/TurboWan2.2-I2V-A14B-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=Wan2_1_Fun_1_3B_InP_SamplingParams,
        pipeline_config_cls=WanI2V480PConfig,
        hf_model_paths=[
            "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam,
        pipeline_config_cls=Wan2_2_TI2V_5B_Config,
        hf_model_paths=[
            "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam,
        pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
        hf_model_paths=[
            "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
            "FastVideo/FastWan2.2-TI2V-5B-Diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=Wan2_2_T2V_A14B_SamplingParam,
        pipeline_config_cls=Wan2_2_T2V_A14B_Config,
        hf_model_paths=[
            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
            "nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4",
        ],
    )
    register_configs(
        sampling_param_cls=Wan2_2_I2V_A14B_SamplingParam,
        pipeline_config_cls=Wan2_2_I2V_A14B_Config,
        hf_model_paths=["Wan-AI/Wan2.2-I2V-A14B-Diffusers"],
    )
    register_configs(
        sampling_param_cls=LingBotWorldSamplingParams,
        pipeline_config_cls=LingBotWorldCausalDMDConfig,
        hf_model_paths=[
            "IPostYellow/lingbot-world-fast-diffusers",
            "robbyant/lingbot-world-fast-diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=LingBotWorldSamplingParams,
        pipeline_config_cls=LingBotWorldV2CausalDMDConfig,
        hf_model_paths=[
            "robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
        ],
    )
    register_configs(
        sampling_param_cls=FastWanT2V480PConfig,
        pipeline_config_cls=FastWan2_1_T2V_480P_Config,
        hf_model_paths=[
            "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
        ],
    )
    # MOVA
    register_configs(
        sampling_param_cls=MOVA_360P_SamplingParams,
        pipeline_config_cls=MOVA360PConfig,
        model_detectors=[
            lambda hf_id: "mova" in hf_id.lower() and "360p" in hf_id.lower()
        ],
    )
    register_configs(
        sampling_param_cls=MOVA_720P_SamplingParams,
        pipeline_config_cls=MOVA720PConfig,
        model_detectors=[
            lambda hf_id: "mova" in hf_id.lower() and "720p" in hf_id.lower()
        ],
    )
    # FLUX
    register_configs(
        sampling_param_cls=FluxSamplingParams,
        pipeline_config_cls=FluxPipelineConfig,
        hf_model_paths=[
            "black-forest-labs/FLUX.1-dev",
        ],
        model_detectors=[lambda hf_id: "flux.1" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=Flux2KleinSamplingParams,
        pipeline_config_cls=Flux2KleinPipelineConfig,
        hf_model_paths=[
            "black-forest-labs/FLUX.2-klein-4B",
            "black-forest-labs/FLUX.2-klein-9B",
        ],
        model_detectors=[
            lambda hf_id: (
                ("flux.2-klein" in hf_id.lower() or "flux2-klein" in hf_id.lower())
                and "base" not in hf_id.lower()
            )
        ],
    )
    register_configs(
        sampling_param_cls=Flux2KleinBaseSamplingParams,
        pipeline_config_cls=Flux2KleinBasePipelineConfig,
        hf_model_paths=[
            "black-forest-labs/FLUX.2-klein-base-4B",
            "black-forest-labs/FLUX.2-klein-base-9B",
        ],
        model_detectors=[
            lambda hf_id: (
                ("flux.2-klein" in hf_id.lower() or "flux2-klein" in hf_id.lower())
                and "base" in hf_id.lower()
            )
        ],
    )
    register_configs(
        sampling_param_cls=Flux2SamplingParams,
        pipeline_config_cls=Flux2PipelineConfig,
        hf_model_paths=[
            "black-forest-labs/FLUX.2-dev",
            "black-forest-labs/FLUX.2-dev-NVFP4",
        ],
        model_detectors=[
            lambda hf_id: "flux.2" in hf_id.lower() and "klein" not in hf_id.lower()
        ],
    )
    register_configs(
        sampling_param_cls=ZImageTurboSamplingParams,
        pipeline_config_cls=ZImagePipelineConfig,
        hf_model_paths=[
            "Tongyi-MAI/Z-Image-Turbo",
        ],
        model_detectors=[lambda hf_id: "z-image-turbo" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=ZImageSamplingParams,
        pipeline_config_cls=ZImagePipelineConfig,
        hf_model_paths=[
            "Tongyi-MAI/Z-Image",
        ],
        model_detectors=[
            lambda hf_id: "z-image" in hf_id.lower() and "turbo" not in hf_id.lower()
        ],
    )
    # Krea-2 (K2)
    register_configs(
        sampling_param_cls=Krea2SamplingParams,
        pipeline_config_cls=Krea2PipelineConfig,
        hf_model_paths=["krea/Krea-2"],
        model_detectors=[lambda hf_id: "krea-2" in hf_id.lower()],
    )
    # Qwen-Image
    register_configs(
        sampling_param_cls=QwenImageSamplingParams,
        pipeline_config_cls=QwenImagePipelineConfig,
        hf_model_paths=["Qwen/Qwen-Image", "nvidia/Qwen-Image-NVFP4"],
        model_detectors=[
            lambda hf_id: (
                "qwen-image" in hf_id.lower()
                and "edit" not in hf_id.lower()
                and "layered" not in hf_id.lower()
                and "2512" not in hf_id.lower()
            )
        ],
    )
    register_configs(
        sampling_param_cls=QwenImage2512SamplingParams,
        pipeline_config_cls=QwenImagePipelineConfig,
        hf_model_paths=["Qwen/Qwen-Image-2512"],
        model_detectors=[lambda hf_id: "qwen-image-2512" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=QwenImageSamplingParams,
        pipeline_config_cls=QwenImageEditPipelineConfig,
        hf_model_paths=["Qwen/Qwen-Image-Edit"],
        model_detectors=[
            lambda hf_id: (
                "qwen-image-edit" in hf_id.lower()
                and "2509" not in hf_id.lower()
                and "2511" not in hf_id.lower()
            )
        ],
    )

    register_configs(
        sampling_param_cls=QwenImageEditPlusSamplingParams,
        pipeline_config_cls=QwenImageEditPlusPipelineConfig,
        hf_model_paths=["Qwen/Qwen-Image-Edit-2509"],
        model_detectors=[lambda hf_id: "qwen-image-edit-2509" in hf_id.lower()],
    )

    register_configs(
        sampling_param_cls=QwenImageEditPlusSamplingParams,
        pipeline_config_cls=QwenImageEditPlus_2511_PipelineConfig,
        hf_model_paths=["Qwen/Qwen-Image-Edit-2511"],
        model_detectors=[lambda hf_id: "qwen-image-edit-2511" in hf_id.lower()],
    )

    register_configs(
        sampling_param_cls=QwenImageLayeredSamplingParams,
        pipeline_config_cls=QwenImageLayeredPipelineConfig,
        hf_model_paths=["Qwen/Qwen-Image-Layered"],
        model_detectors=[lambda hf_id: "qwen-image-layered" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=StableDiffusion3SamplingParams,
        pipeline_config_cls=StableDiffusion3PipelineConfig,
        hf_model_paths=[
            "stabilityai/stable-diffusion-3-medium",
            "stabilityai/stable-diffusion-3-medium-diffusers",
            "stabilityai/stable-diffusion-3.5-medium",
            "stabilityai/stable-diffusion-3.5-medium-diffusers",
            "stabilityai/stable-diffusion-3.5-large",
            "stabilityai/stable-diffusion-3.5-large-diffusers",
        ],
        model_detectors=[
            lambda hf_id: (
                "stable-diffusion-3-medium" in hf_id.lower()
                or "stable-diffusion-3.5-medium" in hf_id.lower()
                or "stable-diffusion-3.5-large" in hf_id.lower()
                or "sd3-medium" in hf_id.lower()
                or "sd3.5-medium" in hf_id.lower()
                or "sd3.5-large" in hf_id.lower()
            )
        ],
    )

    register_configs(
        sampling_param_cls=GlmImageSamplingParams,
        pipeline_config_cls=GlmImagePipelineConfig,
        model_detectors=[lambda hf_id: "glm-image" in hf_id.lower()],
    )
    register_configs(
        sampling_param_cls=Hunyuan3DSamplingParams,
        pipeline_config_cls=Hunyuan3D2PipelineConfig,
        hf_model_paths=[
            "tencent/Hunyuan3D-2",
        ],
        model_detectors=[lambda hf_id: "hunyuan3d" in hf_id.lower()],
    )

    # Helios
    register_configs(
        sampling_param_cls=HeliosT2VSamplingParams,
        pipeline_config_cls=HeliosT2VConfig,
        hf_model_paths=[
            "BestWishYsh/Helios-Base",
        ],
        model_detectors=[
            lambda hf_id: (
                "helios" in hf_id.lower()
                and "mid" not in hf_id.lower()
                and "distill" not in hf_id.lower()
            )
        ],
    )
    register_configs(
        sampling_param_cls=HeliosMidSamplingParams,
        pipeline_config_cls=HeliosMidConfig,
        hf_model_paths=[
            "BestWishYsh/Helios-Mid",
        ],
    )
    register_configs(
        sampling_param_cls=HeliosDistilledSamplingParams,
        pipeline_config_cls=HeliosDistilledConfig,
        hf_model_paths=[
            "BestWishYsh/Helios-Distilled",
        ],
    )

    # SANA-WM (register BEFORE generic SANA T2I to prevent "sana" detector false-match)
    register_configs(
        sampling_param_cls=SanaWMSamplingParams,
        pipeline_config_cls=SanaWMPipelineConfig,
        hf_model_paths=[
            "Efficient-Large-Model/SANA-WM_bidirectional",
            "Efficient-Large-Model/SANA-WM_streaming",
        ],
        model_detectors=[
            # Match "sana-wm" or "sana_wm" but NOT plain T2I "sana" checkpoints.
            lambda hf_id: ("sana-wm" in hf_id.lower() or "sana_wm" in hf_id.lower()),
        ],
    )

    # Cosmos3 — single checkpoint serves T2V, I2V, and T2I. Mode is dispatched
    # per-request inside the pipeline from ``num_frames`` and ``image_path``.
    # Both Nano (8B) and Super (32B) share the same pipeline; arch dimensions
    # come from ``transformer/config.json`` via ``update_model_arch``.
    register_configs(
        sampling_param_cls=Cosmos3SamplingParams,
        pipeline_config_cls=Cosmos3Config,
        hf_model_paths=[
            "nvidia/Cosmos3-Nano",
            "nvidia/Cosmos3-Super",
            "nvidia/Cosmos3-Super-Text2Image",
            "nvidia/Cosmos3-Super-Image2Video",
        ],
        model_detectors=[lambda hf_id: "cosmos3omnidiffuserspipeline" in hf_id.lower()],
    )

    # SANA
    register_configs(
        sampling_param_cls=SanaSamplingParams,
        pipeline_config_cls=SanaPipelineConfig,
        hf_model_paths=[
            "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers",
            "Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers",
            "Efficient-Large-Model/Sana_1600M_1024px_diffusers",
            "Efficient-Large-Model/Sana_600M_1024px_diffusers",
            "Efficient-Large-Model/Sana_1600M_512px_diffusers",
            "Efficient-Large-Model/Sana_600M_512px_diffusers",
        ],
        model_detectors=[
            lambda hf_id: (
                "sana" in hf_id.lower()
                and "sana-wm" not in hf_id.lower()
                and "sana_wm" not in hf_id.lower()
            )
        ],
    )

    # FireRed-Image-Edit
    register_configs(
        sampling_param_cls=QwenImageEditPlusSamplingParams,
        pipeline_config_cls=QwenImageEditPlusPipelineConfig,
        hf_model_paths=[
            "FireRedTeam/FireRed-Image-Edit-1.0",
            "FireRedTeam/FireRed-Image-Edit-1.1",
        ],
    )

    # ErnieImage
    register_configs(
        sampling_param_cls=ErnieImageSamplingParams,
        pipeline_config_cls=ErnieImagePipelineConfig,
        hf_model_paths=[
            "baidu/ERNIE-Image",
            "baidu/ERNIE-Image-Turbo",
        ],
        model_detectors=[
            lambda hf_id: "ernie-image" in hf_id.lower(),
        ],
    )

    # JoyAI
    register_configs(
        sampling_param_cls=JoyImageEditSamplingParams,
        pipeline_config_cls=JoyImageEditPipelineConfig,
        hf_model_paths=[
            "jdopensource/JoyAI-Image-Edit-Diffusers",
        ],
        model_detectors=[
            lambda hf_id: "joyai-image-edit" in hf_id.lower(),
        ],
    )
    register_configs(
        sampling_param_cls=JoyEchoSamplingParams,
        pipeline_config_cls=JoyEchoPipelineConfig,
        hf_model_paths=[
            "jdopensource/JoyAI-Echo",
        ],
        model_detectors=[
            lambda hf_id: ("joy-echo" in hf_id.lower() or "joyai-echo" in hf_id.lower())
            and "image-edit" not in hf_id.lower(),
        ],
    )

    # Ideogram 4
    register_configs(
        sampling_param_cls=Ideogram4SamplingParams,
        pipeline_config_cls=Ideogram4PipelineConfig,
        hf_model_paths=[
            "ideogram-ai/ideogram-4-fp8",
            "ideogram-ai/ideogram-4-nf4",
            "Comfy-Org/Ideogram-4",
        ],
        model_detectors=[
            lambda hf_id: "ideogram4pipeline" in hf_id.lower(),
            lambda hf_id: "ideogram-4-fp8" in hf_id.lower(),
            lambda hf_id: "ideogram-4-nf4" in hf_id.lower(),
            lambda hf_id: "comfy-org/ideogram-4" in hf_id.lower(),
            lambda hf_id: "comfy-org--ideogram-4" in hf_id.lower(),
        ],
    )


_register_configs()


def is_known_non_diffusers_multimodal_model(model_path: str) -> bool:
    model_path_lower = model_path.lower()
    return any(
        pattern in model_path_lower
        for pattern in KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS
    )


def get_non_diffusers_pipeline_name(model_path: str) -> Optional[str]:
    """Get the pipeline name for a known non-diffusers model."""
    model_path_lower = model_path.lower()
    for pattern, pipeline_name in KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS.items():
        if pattern in model_path_lower:
            return pipeline_name
    return None