feat(download): add configurable base model download exclusions

This commit is contained in:
Will Miao
2026-03-26 23:06:14 +08:00
parent 5b065b47d4
commit a5191414cc
22 changed files with 988 additions and 4 deletions

View File

@@ -13,6 +13,7 @@ from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
from ..utils.constants import (
CARD_PREVIEW_WIDTH,
DIFFUSION_MODEL_BASE_MODELS,
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
VALID_LORA_TYPES,
)
from ..utils.civitai_utils import rewrite_preview_url
@@ -228,7 +229,9 @@ class DownloadManager:
# Update status based on result
if task_id in self._active_downloads:
self._active_downloads[task_id]["status"] = (
"completed" if result["success"] else "failed"
result.get("status", "completed")
if result["success"]
else "failed"
)
if not result["success"]:
self._active_downloads[task_id]["error"] = result.get(
@@ -352,10 +355,54 @@ class DownloadManager:
"error": f'Model type "{model_type_from_info}" is not supported for download',
}
excluded_base_models = get_settings_manager().get_download_skip_base_models()
base_model_value = version_info.get("baseModel", "")
if (
isinstance(base_model_value, str)
and base_model_value in SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS
and base_model_value in excluded_base_models
):
file_name = ""
files = version_info.get("files")
if isinstance(files, list):
primary_file = next(
(
file_info
for file_info in files
if isinstance(file_info, dict) and file_info.get("primary")
),
None,
)
selected_file = primary_file
if selected_file is None:
selected_file = next(
(file_info for file_info in files if isinstance(file_info, dict)),
None,
)
if isinstance(selected_file, dict):
raw_file_name = selected_file.get("name", "")
if isinstance(raw_file_name, str):
file_name = raw_file_name.strip()
message = (
f"Skipped download for '{file_name or version_info.get('name') or f'model_version:{model_version_id or model_id}'}' "
f"because base model '{base_model_value}' is excluded in settings"
)
logger.info(message)
return {
"success": True,
"skipped": True,
"status": "skipped",
"reason": "base_model_excluded",
"message": message,
"base_model": base_model_value,
"file_name": file_name,
"download_id": download_id,
}
# Check if this checkpoint should be treated as a diffusion model based on baseModel
is_diffusion_model = False
if model_type == "checkpoint":
base_model_value = version_info.get("baseModel", "")
if base_model_value in DIFFUSION_MODEL_BASE_MODELS:
is_diffusion_model = True
logger.info(

View File

@@ -11,7 +11,11 @@ from typing import Any, Awaitable, Dict, Iterable, List, Mapping, Optional, Sequ
from platformdirs import user_config_dir
from ..utils.constants import DEFAULT_HASH_CHUNK_SIZE_MB, DEFAULT_PRIORITY_TAG_CONFIG
from ..utils.constants import (
DEFAULT_HASH_CHUNK_SIZE_MB,
DEFAULT_PRIORITY_TAG_CONFIG,
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
)
from ..utils.preview_selection import VALID_MATURE_BLUR_LEVELS
from ..utils.settings_paths import APP_NAME, ensure_settings_file, get_legacy_settings_path
from ..utils.tag_priorities import (
@@ -73,6 +77,7 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
"update_flag_strategy": "same_base",
"auto_organize_exclusions": [],
"metadata_refresh_skip_paths": [],
"download_skip_base_models": [],
}
@@ -276,6 +281,21 @@ class SettingsManager:
self.settings["metadata_refresh_skip_paths"] = []
inserted_defaults = True
if "download_skip_base_models" in self.settings:
normalized_skip_base_models = self.normalize_download_skip_base_models(
self.settings.get("download_skip_base_models")
)
if normalized_skip_base_models != self.settings.get(
"download_skip_base_models"
):
self.settings["download_skip_base_models"] = (
normalized_skip_base_models
)
updated_existing = True
else:
self.settings["download_skip_base_models"] = []
inserted_defaults = True
had_mature_level = "mature_blur_level" in self.settings
raw_mature_level = self.settings.get("mature_blur_level")
normalized_mature_level = self.normalize_mature_blur_level(raw_mature_level)
@@ -964,6 +984,45 @@ class SettingsManager:
self._save_settings()
return skip_paths
def normalize_download_skip_base_models(self, value: Any) -> List[str]:
if value is None:
return []
if isinstance(value, str):
candidates: Iterable[str] = (
value.replace("\n", ",").replace(";", ",").split(",")
)
elif isinstance(value, Sequence) and not isinstance(
value, (bytes, bytearray, str)
):
candidates = value
else:
return []
base_models: List[str] = []
seen = set()
for raw in candidates:
if not isinstance(raw, str):
continue
token = raw.strip()
if not token or token not in SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS:
continue
if token in seen:
continue
seen.add(token)
base_models.append(token)
return base_models
def get_download_skip_base_models(self) -> List[str]:
base_models = self.normalize_download_skip_base_models(
self.settings.get("download_skip_base_models")
)
if base_models != self.settings.get("download_skip_base_models"):
self.settings["download_skip_base_models"] = base_models
self._save_settings()
return base_models
def get_extra_folder_paths(self) -> Dict[str, List[str]]:
"""Get extra folder paths for the active library.
@@ -1032,6 +1091,8 @@ class SettingsManager:
value = self.normalize_auto_organize_exclusions(value)
elif key == "metadata_refresh_skip_paths":
value = self.normalize_metadata_refresh_skip_paths(value)
elif key == "download_skip_base_models":
value = self.normalize_download_skip_base_models(value)
elif key == "mature_blur_level":
value = self.normalize_mature_blur_level(value)
self.settings[key] = value

View File

@@ -113,3 +113,59 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset(
"Qwen",
]
)
# Supported baseModel values for download exclusion settings.
# Keep this aligned with static/js/utils/constants.js, excluding the generic "Other" value.
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS = frozenset(
[
"SD 1.4",
"SD 1.5",
"SD 1.5 LCM",
"SD 1.5 Hyper",
"SD 2.0",
"SD 2.1",
"SD 3",
"SD 3.5",
"SD 3.5 Medium",
"SD 3.5 Large",
"SD 3.5 Large Turbo",
"SDXL 1.0",
"SDXL Lightning",
"SDXL Hyper",
"Flux.1 D",
"Flux.1 S",
"Flux.1 Krea",
"Flux.1 Kontext",
"Flux.2 D",
"Flux.2 Klein 9B",
"Flux.2 Klein 9B-base",
"Flux.2 Klein 4B",
"Flux.2 Klein 4B-base",
"AuraFlow",
"Chroma",
"PixArt a",
"PixArt E",
"Hunyuan 1",
"Lumina",
"Kolors",
"NoobAI",
"Illustrious",
"Pony",
"HiDream",
"Qwen",
"ZImageTurbo",
"ZImageBase",
"SVD",
"LTXV",
"LTXV2",
"Wan Video",
"Wan Video 1.3B t2v",
"Wan Video 14B t2v",
"Wan Video 14B i2v 480p",
"Wan Video 14B i2v 720p",
"Wan Video 2.2 TI2V-5B",
"Wan Video 2.2 T2V-A14B",
"Wan Video 2.2 I2V-A14B",
"Hunyuan Video",
]
)