registry ¶

def get_default_preset(model_path: str) -> str | None:
    """Return the ``default_preset`` name for a model path."""
    config_info = _get_config_info(model_path, raise_on_missing=False)
    if config_info is None:
        return None
    return config_info.default_preset

def get_model_family(model_path: str) -> str | None:
    """Return the ``model_family`` string for a model path, or ``None``."""
    config_info = _get_config_info(model_path, raise_on_missing=False)
    if config_info is None:
        return None
    return config_info.model_family

def get_preset_selection(model_path: str) -> tuple[str | None, str | None]:
    """Return ``(default_preset, model_family)`` for a model path.

    Single-lookup variant of :func:`get_default_preset` +
    :func:`get_model_family`; callers that need both should prefer this
    to avoid walking the registry twice.
    """
    config_info = _get_config_info(model_path, raise_on_missing=False)
    if config_info is None:
        return None, None
    return config_info.default_preset, config_info.model_family

def get_registered_model_paths() -> list[str]:
    """Return all registered HuggingFace model paths.

    Useful for UIs and tooling that need to enumerate supported models.
    """
    return sorted(_MODEL_HF_PATH_TO_NAME.keys())

def get_registered_models_with_workloads(workload_type: str | None = None, ) -> list[dict[str, Any]]:
    """Return models with workload metadata, optionally filtered by workload.

    Args:
        workload_type: If set (e.g. "t2v", "i2v", "t2i"), only return models
            that support this workload. If None, return all with workload_types.

    Returns:
        List of dicts with keys: id, label, workload_types.
    """
    result: list[dict[str, Any]] = []
    for path in sorted(_MODEL_HF_PATH_TO_NAME.keys()):
        model_id = _MODEL_HF_PATH_TO_NAME[path]
        config_info = _CONFIG_REGISTRY.get(model_id)
        if config_info is None:
            continue
        workload_values = [w.value for w in config_info.workload_types]
        if workload_type is not None and workload_type.lower() not in workload_values:
            continue
        label = path.split("/")[-1].replace("-", " ").replace("_", " ")
        result.append({
            "id": path,
            "label": label,
            "workload_types": workload_values,
        })
    return result

def register_configs(
    sampling_param_cls: type[SamplingParam] | None,
    pipeline_config_cls: type[PipelineConfig],
    workload_types: tuple[WorkloadType, ...],
    hf_model_paths: list[str] | None = None,
    model_detectors: list[Callable[[str], bool]] | None = None,
    model_family: str | None = None,
    default_preset: str | None = None,
    pipeline_cls_name: str | None = None,
) -> None:
    """Register config classes for a model family.

    workload_types declares which UI workload options this config supports.
    Use () for configs not exposed as workload options.
    """
    model_id = str(len(_CONFIG_REGISTRY))

    _CONFIG_REGISTRY[model_id] = ConfigInfo(
        sampling_param_cls=sampling_param_cls,
        pipeline_config_cls=pipeline_config_cls,
        workload_types=workload_types,
        model_family=model_family,
        default_preset=default_preset,
        pipeline_cls_name=pipeline_cls_name,
    )

    if hf_model_paths:
        for path in hf_model_paths:
            if path in _MODEL_HF_PATH_TO_NAME:
                logger.warning("Model path '%s' is already mapped to '%s' and will be overwritten by '%s'.", path,
                               _MODEL_HF_PATH_TO_NAME[path], 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))