alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-25 15:15:44 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(download): auto-route diffusion models to unet folder based on baseModel, see #770

Browse Source 
CivitAI does not distinguish between checkpoint and diffusion model types -
both are labeled as "checkpoint". For certain base model types like
"ZImageTurbo", all models are actually diffusion models and should be
saved to the unet/diffusion model folder instead of the checkpoint folder.

- Add DIFFUSION_MODEL_BASE_MODELS constant for known diffusion model types
- Add default_unet_root setting with auto-set logic
- Route downloads to unet folder when baseModel matches known diffusion types

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Will Miao
2026-01-18 11:57:23 +08:00
parent aab1797269
commit dad549f65f
3 changed files with 64 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
py/services/download_manager.py
View File

@@ -10,7 +10,7 @@ import uuid
from typing import Dict, List, Optional, Set, Tuple
from urllib.parse import urlparse
from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
from ..utils.constants import CARD_PREVIEW_WIDTH, VALID_LORA_TYPES
from ..utils.constants import CARD_PREVIEW_WIDTH, DIFFUSION_MODEL_BASE_MODELS, VALID_LORA_TYPES
from ..utils.civitai_utils import rewrite_preview_url
from ..utils.preview_selection import select_preview_media
from ..utils.utils import sanitize_folder_name
@@ -343,6 +343,14 @@ class DownloadManager:
"error": f'Model type "{model_type_from_info}" is not supported for download',
}
# 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(f"baseModel '{base_model_value}' is a known diffusion model, routing to unet folder")
# Case 2: model_version_id was None, check after getting version_info
if model_version_id is None:
version_id = version_info.get("id")
@@ -377,11 +385,16 @@ class DownloadManager:
settings_manager = get_settings_manager()
# Set save_dir based on model type
if model_type == "checkpoint":
default_path = settings_manager.get("default_checkpoint_root")
if is_diffusion_model:
default_path = settings_manager.get("default_unet_root")
error_msg = "Default unet root path not set in settings"
else:
default_path = settings_manager.get("default_checkpoint_root")
error_msg = "Default checkpoint root path not set in settings"
if not default_path:
return {
"success": False,
"error": "Default checkpoint root path not set in settings",
"error": error_msg,
}
save_dir = default_path
elif model_type == "lora":