Files
wehub-resource-sync cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:22:06 +08:00

588 lines
28 KiB
Python

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import (
Union,
)
from pydantic import Discriminator, TypeAdapter, ValidationError
from typing_extensions import Annotated, Any
from invokeai.app.services.config.config_default import get_config
from invokeai.app.util.misc import uuid_string
from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS
from invokeai.backend.model_manager.configs.base import Config_Base
from invokeai.backend.model_manager.configs.clip_embed import CLIPEmbed_Diffusers_G_Config, CLIPEmbed_Diffusers_L_Config
from invokeai.backend.model_manager.configs.clip_vision import CLIPVision_Diffusers_Config
from invokeai.backend.model_manager.configs.controlnet import (
ControlAdapterDefaultSettings,
ControlNet_Checkpoint_Anima_Config,
ControlNet_Checkpoint_FLUX_Config,
ControlNet_Checkpoint_SD1_Config,
ControlNet_Checkpoint_SD2_Config,
ControlNet_Checkpoint_SDXL_Config,
ControlNet_Checkpoint_ZImage_Config,
ControlNet_Diffusers_FLUX_Config,
ControlNet_Diffusers_SD1_Config,
ControlNet_Diffusers_SD2_Config,
ControlNet_Diffusers_SDXL_Config,
)
from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig
from invokeai.backend.model_manager.configs.flux_redux import FLUXRedux_Checkpoint_Config
from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError
from invokeai.backend.model_manager.configs.ip_adapter import (
IPAdapter_Checkpoint_FLUX_Config,
IPAdapter_Checkpoint_SD1_Config,
IPAdapter_Checkpoint_SD2_Config,
IPAdapter_Checkpoint_SDXL_Config,
IPAdapter_InvokeAI_SD1_Config,
IPAdapter_InvokeAI_SD2_Config,
IPAdapter_InvokeAI_SDXL_Config,
)
from invokeai.backend.model_manager.configs.llava_onevision import LlavaOnevision_Diffusers_Config
from invokeai.backend.model_manager.configs.lora import (
ControlLoRA_LyCORIS_FLUX_Config,
LoRA_Diffusers_Flux2_Config,
LoRA_Diffusers_FLUX_Config,
LoRA_Diffusers_SD1_Config,
LoRA_Diffusers_SD2_Config,
LoRA_Diffusers_SDXL_Config,
LoRA_Diffusers_ZImage_Config,
LoRA_LyCORIS_Anima_Config,
LoRA_LyCORIS_Flux2_Config,
LoRA_LyCORIS_FLUX_Config,
LoRA_LyCORIS_QwenImage_Config,
LoRA_LyCORIS_SD1_Config,
LoRA_LyCORIS_SD2_Config,
LoRA_LyCORIS_SDXL_Config,
LoRA_LyCORIS_ZImage_Config,
LoRA_OMI_FLUX_Config,
LoRA_OMI_SDXL_Config,
LoraModelDefaultSettings,
)
from invokeai.backend.model_manager.configs.main import (
Main_BnBNF4_FLUX_Config,
Main_Checkpoint_Anima_Config,
Main_Checkpoint_Flux2_Config,
Main_Checkpoint_FLUX_Config,
Main_Checkpoint_QwenImage_Config,
Main_Checkpoint_SD1_Config,
Main_Checkpoint_SD2_Config,
Main_Checkpoint_SDXL_Config,
Main_Checkpoint_SDXLRefiner_Config,
Main_Checkpoint_ZImage_Config,
Main_Diffusers_CogView4_Config,
Main_Diffusers_Flux2_Config,
Main_Diffusers_FLUX_Config,
Main_Diffusers_QwenImage_Config,
Main_Diffusers_SD1_Config,
Main_Diffusers_SD2_Config,
Main_Diffusers_SD3_Config,
Main_Diffusers_SDXL_Config,
Main_Diffusers_SDXLRefiner_Config,
Main_Diffusers_ZImage_Config,
Main_GGUF_Flux2_Config,
Main_GGUF_FLUX_Config,
Main_GGUF_QwenImage_Config,
Main_GGUF_ZImage_Config,
MainModelDefaultSettings,
)
from invokeai.backend.model_manager.configs.qwen3_encoder import (
Qwen3Encoder_Checkpoint_Config,
Qwen3Encoder_GGUF_Config,
Qwen3Encoder_Qwen3Encoder_Config,
)
from invokeai.backend.model_manager.configs.qwen_vl_encoder import (
QwenVLEncoder_Checkpoint_Config,
QwenVLEncoder_Diffusers_Config,
)
from invokeai.backend.model_manager.configs.siglip import SigLIP_Diffusers_Config
from invokeai.backend.model_manager.configs.spandrel import Spandrel_Checkpoint_Config
from invokeai.backend.model_manager.configs.t2i_adapter import (
T2IAdapter_Diffusers_SD1_Config,
T2IAdapter_Diffusers_SDXL_Config,
)
from invokeai.backend.model_manager.configs.t5_encoder import T5Encoder_BnBLLMint8_Config, T5Encoder_T5Encoder_Config
from invokeai.backend.model_manager.configs.text_llm import TextLLM_Diffusers_Config
from invokeai.backend.model_manager.configs.textual_inversion import (
TI_File_SD1_Config,
TI_File_SD2_Config,
TI_File_SDXL_Config,
TI_Folder_SD1_Config,
TI_Folder_SD2_Config,
TI_Folder_SDXL_Config,
)
from invokeai.backend.model_manager.configs.unknown import Unknown_Config
from invokeai.backend.model_manager.configs.vae import (
VAE_Checkpoint_Anima_Config,
VAE_Checkpoint_Flux2_Config,
VAE_Checkpoint_FLUX_Config,
VAE_Checkpoint_QwenImage_Config,
VAE_Checkpoint_SD1_Config,
VAE_Checkpoint_SD2_Config,
VAE_Checkpoint_SDXL_Config,
VAE_Diffusers_Flux2_Config,
VAE_Diffusers_SD1_Config,
VAE_Diffusers_SDXL_Config,
)
from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
from invokeai.backend.model_manager.taxonomy import (
BaseModelType,
ModelFormat,
ModelSourceType,
ModelType,
variant_type_adapter,
)
logger = logging.getLogger(__name__)
app_config = get_config()
# Known model file extensions for sanity checking
_MODEL_EXTENSIONS = {
".safetensors",
".ckpt",
".pt",
".pth",
".bin",
".gguf",
".onnx",
}
# Known config file names for diffusers/transformers models
_CONFIG_FILES = {
"model_index.json",
"config.json",
}
# Maximum number of files in a directory to be considered a model
_MAX_FILES_IN_MODEL_DIR = 50
# Maximum depth to search for model files in directories
_MAX_SEARCH_DEPTH = 2
# The types are listed explicitly because IDEs/LSPs can't identify the correct types
# when AnyModelConfig is constructed dynamically using ModelConfigBase.all_config_classes
AnyModelConfig = Annotated[
Union[
# Main (Pipeline) - diffusers format
Annotated[Main_Diffusers_SD1_Config, Main_Diffusers_SD1_Config.get_tag()],
Annotated[Main_Diffusers_SD2_Config, Main_Diffusers_SD2_Config.get_tag()],
Annotated[Main_Diffusers_SDXL_Config, Main_Diffusers_SDXL_Config.get_tag()],
Annotated[Main_Diffusers_SDXLRefiner_Config, Main_Diffusers_SDXLRefiner_Config.get_tag()],
Annotated[Main_Diffusers_SD3_Config, Main_Diffusers_SD3_Config.get_tag()],
Annotated[Main_Diffusers_FLUX_Config, Main_Diffusers_FLUX_Config.get_tag()],
Annotated[Main_Diffusers_Flux2_Config, Main_Diffusers_Flux2_Config.get_tag()],
Annotated[Main_Diffusers_CogView4_Config, Main_Diffusers_CogView4_Config.get_tag()],
Annotated[Main_Diffusers_QwenImage_Config, Main_Diffusers_QwenImage_Config.get_tag()],
Annotated[Main_Diffusers_ZImage_Config, Main_Diffusers_ZImage_Config.get_tag()],
# Main (Pipeline) - checkpoint format
# IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
# that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
Annotated[Main_Checkpoint_SD1_Config, Main_Checkpoint_SD1_Config.get_tag()],
Annotated[Main_Checkpoint_SD2_Config, Main_Checkpoint_SD2_Config.get_tag()],
Annotated[Main_Checkpoint_SDXL_Config, Main_Checkpoint_SDXL_Config.get_tag()],
Annotated[Main_Checkpoint_SDXLRefiner_Config, Main_Checkpoint_SDXLRefiner_Config.get_tag()],
Annotated[Main_Checkpoint_Flux2_Config, Main_Checkpoint_Flux2_Config.get_tag()],
Annotated[Main_Checkpoint_FLUX_Config, Main_Checkpoint_FLUX_Config.get_tag()],
Annotated[Main_Checkpoint_QwenImage_Config, Main_Checkpoint_QwenImage_Config.get_tag()],
Annotated[Main_Checkpoint_ZImage_Config, Main_Checkpoint_ZImage_Config.get_tag()],
Annotated[Main_Checkpoint_Anima_Config, Main_Checkpoint_Anima_Config.get_tag()],
# Main (Pipeline) - quantized formats
# IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
# that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
Annotated[Main_BnBNF4_FLUX_Config, Main_BnBNF4_FLUX_Config.get_tag()],
Annotated[Main_GGUF_Flux2_Config, Main_GGUF_Flux2_Config.get_tag()],
Annotated[Main_GGUF_FLUX_Config, Main_GGUF_FLUX_Config.get_tag()],
Annotated[Main_GGUF_QwenImage_Config, Main_GGUF_QwenImage_Config.get_tag()],
Annotated[Main_GGUF_ZImage_Config, Main_GGUF_ZImage_Config.get_tag()],
# VAE - checkpoint format
Annotated[VAE_Checkpoint_SD1_Config, VAE_Checkpoint_SD1_Config.get_tag()],
Annotated[VAE_Checkpoint_SD2_Config, VAE_Checkpoint_SD2_Config.get_tag()],
Annotated[VAE_Checkpoint_SDXL_Config, VAE_Checkpoint_SDXL_Config.get_tag()],
Annotated[VAE_Checkpoint_FLUX_Config, VAE_Checkpoint_FLUX_Config.get_tag()],
Annotated[VAE_Checkpoint_Flux2_Config, VAE_Checkpoint_Flux2_Config.get_tag()],
Annotated[VAE_Checkpoint_QwenImage_Config, VAE_Checkpoint_QwenImage_Config.get_tag()],
Annotated[VAE_Checkpoint_Anima_Config, VAE_Checkpoint_Anima_Config.get_tag()],
# VAE - diffusers format
Annotated[VAE_Diffusers_SD1_Config, VAE_Diffusers_SD1_Config.get_tag()],
Annotated[VAE_Diffusers_SDXL_Config, VAE_Diffusers_SDXL_Config.get_tag()],
Annotated[VAE_Diffusers_Flux2_Config, VAE_Diffusers_Flux2_Config.get_tag()],
# ControlNet - checkpoint format
Annotated[ControlNet_Checkpoint_SD1_Config, ControlNet_Checkpoint_SD1_Config.get_tag()],
Annotated[ControlNet_Checkpoint_SD2_Config, ControlNet_Checkpoint_SD2_Config.get_tag()],
Annotated[ControlNet_Checkpoint_SDXL_Config, ControlNet_Checkpoint_SDXL_Config.get_tag()],
Annotated[ControlNet_Checkpoint_FLUX_Config, ControlNet_Checkpoint_FLUX_Config.get_tag()],
Annotated[ControlNet_Checkpoint_ZImage_Config, ControlNet_Checkpoint_ZImage_Config.get_tag()],
Annotated[ControlNet_Checkpoint_Anima_Config, ControlNet_Checkpoint_Anima_Config.get_tag()],
# ControlNet - diffusers format
Annotated[ControlNet_Diffusers_SD1_Config, ControlNet_Diffusers_SD1_Config.get_tag()],
Annotated[ControlNet_Diffusers_SD2_Config, ControlNet_Diffusers_SD2_Config.get_tag()],
Annotated[ControlNet_Diffusers_SDXL_Config, ControlNet_Diffusers_SDXL_Config.get_tag()],
Annotated[ControlNet_Diffusers_FLUX_Config, ControlNet_Diffusers_FLUX_Config.get_tag()],
# LoRA - LyCORIS format
# IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
# that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
Annotated[LoRA_LyCORIS_SD1_Config, LoRA_LyCORIS_SD1_Config.get_tag()],
Annotated[LoRA_LyCORIS_SD2_Config, LoRA_LyCORIS_SD2_Config.get_tag()],
Annotated[LoRA_LyCORIS_SDXL_Config, LoRA_LyCORIS_SDXL_Config.get_tag()],
Annotated[LoRA_LyCORIS_Flux2_Config, LoRA_LyCORIS_Flux2_Config.get_tag()],
Annotated[LoRA_LyCORIS_FLUX_Config, LoRA_LyCORIS_FLUX_Config.get_tag()],
Annotated[LoRA_LyCORIS_ZImage_Config, LoRA_LyCORIS_ZImage_Config.get_tag()],
Annotated[LoRA_LyCORIS_QwenImage_Config, LoRA_LyCORIS_QwenImage_Config.get_tag()],
Annotated[LoRA_LyCORIS_Anima_Config, LoRA_LyCORIS_Anima_Config.get_tag()],
# LoRA - OMI format
Annotated[LoRA_OMI_SDXL_Config, LoRA_OMI_SDXL_Config.get_tag()],
Annotated[LoRA_OMI_FLUX_Config, LoRA_OMI_FLUX_Config.get_tag()],
# LoRA - diffusers format
# IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
# that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
Annotated[LoRA_Diffusers_SD1_Config, LoRA_Diffusers_SD1_Config.get_tag()],
Annotated[LoRA_Diffusers_SD2_Config, LoRA_Diffusers_SD2_Config.get_tag()],
Annotated[LoRA_Diffusers_SDXL_Config, LoRA_Diffusers_SDXL_Config.get_tag()],
Annotated[LoRA_Diffusers_Flux2_Config, LoRA_Diffusers_Flux2_Config.get_tag()],
Annotated[LoRA_Diffusers_FLUX_Config, LoRA_Diffusers_FLUX_Config.get_tag()],
Annotated[LoRA_Diffusers_ZImage_Config, LoRA_Diffusers_ZImage_Config.get_tag()],
# ControlLoRA - diffusers format
Annotated[ControlLoRA_LyCORIS_FLUX_Config, ControlLoRA_LyCORIS_FLUX_Config.get_tag()],
# T5 Encoder - all formats
Annotated[T5Encoder_T5Encoder_Config, T5Encoder_T5Encoder_Config.get_tag()],
Annotated[T5Encoder_BnBLLMint8_Config, T5Encoder_BnBLLMint8_Config.get_tag()],
# Qwen3 Encoder
Annotated[Qwen3Encoder_Qwen3Encoder_Config, Qwen3Encoder_Qwen3Encoder_Config.get_tag()],
Annotated[Qwen3Encoder_Checkpoint_Config, Qwen3Encoder_Checkpoint_Config.get_tag()],
Annotated[Qwen3Encoder_GGUF_Config, Qwen3Encoder_GGUF_Config.get_tag()],
# Qwen VL Encoder (Qwen2.5-VL multimodal encoder for Qwen Image)
Annotated[QwenVLEncoder_Diffusers_Config, QwenVLEncoder_Diffusers_Config.get_tag()],
Annotated[QwenVLEncoder_Checkpoint_Config, QwenVLEncoder_Checkpoint_Config.get_tag()],
# TI - file format
Annotated[TI_File_SD1_Config, TI_File_SD1_Config.get_tag()],
Annotated[TI_File_SD2_Config, TI_File_SD2_Config.get_tag()],
Annotated[TI_File_SDXL_Config, TI_File_SDXL_Config.get_tag()],
# TI - folder format
Annotated[TI_Folder_SD1_Config, TI_Folder_SD1_Config.get_tag()],
Annotated[TI_Folder_SD2_Config, TI_Folder_SD2_Config.get_tag()],
Annotated[TI_Folder_SDXL_Config, TI_Folder_SDXL_Config.get_tag()],
# IP Adapter - InvokeAI format
Annotated[IPAdapter_InvokeAI_SD1_Config, IPAdapter_InvokeAI_SD1_Config.get_tag()],
Annotated[IPAdapter_InvokeAI_SD2_Config, IPAdapter_InvokeAI_SD2_Config.get_tag()],
Annotated[IPAdapter_InvokeAI_SDXL_Config, IPAdapter_InvokeAI_SDXL_Config.get_tag()],
# IP Adapter - checkpoint format
Annotated[IPAdapter_Checkpoint_SD1_Config, IPAdapter_Checkpoint_SD1_Config.get_tag()],
Annotated[IPAdapter_Checkpoint_SD2_Config, IPAdapter_Checkpoint_SD2_Config.get_tag()],
Annotated[IPAdapter_Checkpoint_SDXL_Config, IPAdapter_Checkpoint_SDXL_Config.get_tag()],
Annotated[IPAdapter_Checkpoint_FLUX_Config, IPAdapter_Checkpoint_FLUX_Config.get_tag()],
# T2I Adapter - diffusers format
Annotated[T2IAdapter_Diffusers_SD1_Config, T2IAdapter_Diffusers_SD1_Config.get_tag()],
Annotated[T2IAdapter_Diffusers_SDXL_Config, T2IAdapter_Diffusers_SDXL_Config.get_tag()],
# Misc models
Annotated[Spandrel_Checkpoint_Config, Spandrel_Checkpoint_Config.get_tag()],
Annotated[CLIPEmbed_Diffusers_G_Config, CLIPEmbed_Diffusers_G_Config.get_tag()],
Annotated[CLIPEmbed_Diffusers_L_Config, CLIPEmbed_Diffusers_L_Config.get_tag()],
Annotated[CLIPVision_Diffusers_Config, CLIPVision_Diffusers_Config.get_tag()],
Annotated[SigLIP_Diffusers_Config, SigLIP_Diffusers_Config.get_tag()],
Annotated[FLUXRedux_Checkpoint_Config, FLUXRedux_Checkpoint_Config.get_tag()],
Annotated[LlavaOnevision_Diffusers_Config, LlavaOnevision_Diffusers_Config.get_tag()],
Annotated[TextLLM_Diffusers_Config, TextLLM_Diffusers_Config.get_tag()],
Annotated[ExternalApiModelConfig, ExternalApiModelConfig.get_tag()],
# Unknown model (fallback)
Annotated[Unknown_Config, Unknown_Config.get_tag()],
],
Discriminator(Config_Base.get_model_discriminator_value),
]
AnyModelConfigValidator = TypeAdapter[AnyModelConfig](AnyModelConfig)
"""Pydantic TypeAdapter for the AnyModelConfig union, used for parsing and validation.
If you need to parse/validate a dict or JSON into an AnyModelConfig, you should probably use
ModelConfigFactory.from_dict or ModelConfigFactory.from_json instead as they may implement
additional logic in the future.
"""
@dataclass
class ModelClassificationResult:
"""Result of attempting to classify a model on disk into a specific model config.
Attributes:
match: The best matching model config, or None if no match was found.
results: A mapping of model config class names to either an instance of that class (if it matched)
or an Exception (if it didn't match or an error occurred during matching).
"""
config: AnyModelConfig | None
details: dict[str, AnyModelConfig | Exception]
@property
def all_matches(self) -> list[AnyModelConfig]:
"""Returns a list of all matching model configs found."""
return [r for r in self.details.values() if isinstance(r, Config_Base)]
@property
def match_count(self) -> int:
"""Returns the number of matching model configs found."""
return len(self.all_matches)
class ModelConfigFactory:
@staticmethod
def from_dict(fields: dict[str, Any]) -> AnyModelConfig:
"""Return the appropriate config object from raw dict values."""
model = AnyModelConfigValidator.validate_python(fields)
return model
@staticmethod
def from_json(json: str | bytes | bytearray) -> AnyModelConfig:
"""Return the appropriate config object from json."""
model = AnyModelConfigValidator.validate_json(json)
return model
@staticmethod
def build_common_fields(
mod: ModelOnDisk,
override_fields: dict[str, Any] | None = None,
) -> dict[str, Any]:
"""Builds the common fields for all model configs.
Args:
mod: The model on disk to extract fields from.
overrides: A optional dictionary of fields to override. These fields will take precedence over the values
extracted from the model on disk.
- Casts string fields to their Enum types.
- Does not validate the fields against the model config schema.
"""
_overrides: dict[str, Any] = override_fields or {}
fields: dict[str, Any] = {}
if "type" in _overrides:
fields["type"] = ModelType(_overrides["type"])
if "format" in _overrides:
fields["format"] = ModelFormat(_overrides["format"])
if "base" in _overrides:
fields["base"] = BaseModelType(_overrides["base"])
if "source_type" in _overrides:
fields["source_type"] = ModelSourceType(_overrides["source_type"])
if "variant" in _overrides:
fields["variant"] = variant_type_adapter.validate_strings(_overrides["variant"])
fields["path"] = mod.path.as_posix()
fields["source"] = _overrides.get("source") or fields["path"]
fields["source_type"] = _overrides.get("source_type") or ModelSourceType.Path
fields["name"] = _overrides.get("name") or mod.name
fields["hash"] = _overrides.get("hash") or mod.hash()
fields["key"] = _overrides.get("key") or uuid_string()
fields["description"] = _overrides.get("description")
fields["file_size"] = _overrides.get("file_size") or mod.size()
return fields
@staticmethod
def _validate_path_looks_like_model(path: Path) -> None:
"""Perform basic sanity checks to ensure a path looks like a model.
This prevents wasting time trying to identify obviously non-model paths like
home directories or downloads folders. Raises RuntimeError if the path doesn't
pass basic checks.
Args:
path: The path to validate
Raises:
ValueError: If the path doesn't look like a model
"""
if path.is_file():
# For files, just check the extension
if path.suffix.lower() not in _MODEL_EXTENSIONS:
raise ValueError(
f"File extension {path.suffix} is not a recognized model format. "
f"Expected one of: {', '.join(sorted(_MODEL_EXTENSIONS))}"
)
else:
# For directories, do a quick file count check with early exit
total_files = 0
# Ignore hidden files and directories
paths_to_check = (
p
for p in path.rglob("*")
if not p.name.startswith(".") and not any(part.startswith(".") for part in p.parts)
)
for item in paths_to_check:
if item.is_file():
total_files += 1
if total_files > _MAX_FILES_IN_MODEL_DIR:
raise ValueError(
f"Directory contains more than {_MAX_FILES_IN_MODEL_DIR} files. "
"This looks like a general-purpose directory rather than a model. "
"Please provide a path to a specific model file or model directory."
)
# Check if it has config files at root (diffusers/transformers marker)
has_root_config = any((path / config).exists() for config in _CONFIG_FILES)
if has_root_config:
# Has a config file, looks like a valid model directory
return
# Otherwise, search for model files within depth limit
def find_model_files(current_path: Path, depth: int) -> bool:
if depth > _MAX_SEARCH_DEPTH:
return False
try:
for item in current_path.iterdir():
if item.is_file() and item.suffix.lower() in _MODEL_EXTENSIONS:
return True
elif item.is_dir() and find_model_files(item, depth + 1):
return True
except PermissionError:
pass
return False
if not find_model_files(path, 0):
raise ValueError(
f"No model files or config files found in directory {path}. "
f"Expected to find model files with extensions: {', '.join(sorted(_MODEL_EXTENSIONS))} "
f"or config files: {', '.join(sorted(_CONFIG_FILES))}"
)
@staticmethod
def matches_sort_key(m: AnyModelConfig) -> int:
"""Sort key function to prioritize model config matches in case of multiple matches."""
# It is possible that we have multiple matches. We need to prioritize them.
# Known cases where multiple matches can occur:
# - SD main models can look like a LoRA when they have merged in LoRA weights. Prefer the main model.
# - SD main models in diffusers format can look like a CLIP Embed; they have a text_encoder folder with
# a config.json file. Prefer the main model.
# Given the above cases, we can prioritize the matches by type. If we find more cases, we may need a more
# sophisticated approach.
match m.type:
case ModelType.Main:
return 0
case ModelType.LoRA:
return 1
case ModelType.CLIPEmbed:
return 2
case _:
return 3
@staticmethod
def from_model_on_disk(
mod: str | Path | ModelOnDisk,
override_fields: dict[str, Any] | None = None,
hash_algo: HASHING_ALGORITHMS = "blake3_single",
allow_unknown: bool = True,
) -> ModelClassificationResult:
"""Classify a model on disk and return the best matching model config.
Args:
mod: The model on disk to classify. Can be a path (str or Path) or a ModelOnDisk instance.
override_fields: Optional dictionary of fields to override. These fields will take precedence
over the values extracted from the model on disk, but this cannot force a match if the
model on disk doesn't actually match the config class.
hash_algo: The hashing algorithm to use when computing the model hash if needed.
Returns:
A ModelClassificationResult containing the best matching model config (or None if no match)
and a mapping of all attempted model config classes to either an instance of that class (if it matched)
or an Exception (if it didn't match or an error occurred during matching).
Raises:
ValueError: If the provided path doesn't look like a model.
"""
if isinstance(mod, Path | str):
mod = ModelOnDisk(Path(mod), hash_algo)
# Perform basic sanity checks before attempting any config matching
# This rejects obviously non-model paths early, saving time
ModelConfigFactory._validate_path_looks_like_model(mod.path)
# We will always need these fields to build any model config.
fields = ModelConfigFactory.build_common_fields(mod, override_fields)
# Store results as a mapping of config class to either an instance of that class or an exception
# that was raised when trying to build it.
details: dict[str, AnyModelConfig | Exception] = {}
# Try to build an instance of each model config class that uses the classify API.
# Each class will either return an instance of itself or raise NotAMatch if it doesn't match.
# Other exceptions may be raised if something unexpected happens during matching or building.
for candidate_class in filter(lambda x: x is not Unknown_Config, Config_Base.CONFIG_CLASSES):
candidate_name = candidate_class.__name__
try:
# Technically, from_model_on_disk returns a Config_Base, but in practice it will always be a member of
# the AnyModelConfig union.
details[candidate_name] = candidate_class.from_model_on_disk(mod, fields) # type: ignore
except NotAMatchError as e:
# This means the model didn't match this config class. It's not an error, just no match.
details[candidate_name] = e
except ValidationError as e:
# This means the model matched, but we couldn't create the pydantic model instance for the config.
# Maybe invalid overrides were provided?
details[candidate_name] = e
except Exception as e:
# Some other unexpected error occurred. Store the exception for reporting later.
details[candidate_name] = e
# Extract just the successful matches
matches = [r for r in details.values() if isinstance(r, Config_Base)]
if not matches:
if not allow_unknown:
# No matches and we are not allowed to fall back to Unknown_Config
return ModelClassificationResult(config=None, details=details)
else:
# Fall back to Unknown_Config
# This should always succeed as Unknown_Config.from_model_on_disk never raises NotAMatch
config = Unknown_Config.from_model_on_disk(mod, fields)
details[Unknown_Config.__name__] = config
return ModelClassificationResult(config=config, details=details)
matches.sort(key=ModelConfigFactory.matches_sort_key)
config = matches[0]
# Now do any post-processing needed for specific model types/bases/etc.
match config.type:
case ModelType.Main:
# Pass variant if available (e.g., for Flux2 models)
variant = getattr(config, "variant", None)
config.default_settings = MainModelDefaultSettings.from_base(config.base, variant)
case ModelType.ControlNet | ModelType.T2IAdapter | ModelType.ControlLoRa:
config.default_settings = ControlAdapterDefaultSettings.from_model_name(config.name)
case ModelType.LoRA:
config.default_settings = LoraModelDefaultSettings()
case _:
pass
return ModelClassificationResult(config=config, details=details)
MODEL_NAME_TO_PREPROCESSOR = {
"canny": "canny_image_processor",
"mlsd": "mlsd_image_processor",
"depth": "depth_anything_image_processor",
"bae": "normalbae_image_processor",
"normal": "normalbae_image_processor",
"sketch": "pidi_image_processor",
"scribble": "lineart_image_processor",
"lineart anime": "lineart_anime_image_processor",
"lineart_anime": "lineart_anime_image_processor",
"lineart": "lineart_image_processor",
"soft": "hed_image_processor",
"softedge": "hed_image_processor",
"hed": "hed_image_processor",
"shuffle": "content_shuffle_image_processor",
"pose": "dw_openpose_image_processor",
"mediapipe": "mediapipe_face_processor",
"pidi": "pidi_image_processor",
"zoe": "zoe_depth_image_processor",
"color": "color_map_image_processor",
}