mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 11:11:26 -03:00
fix(download): align location-step root selection with backend diffusion routing
The download modal's location step decided between checkpoint and unet roots using only the CivitAI file-type signal, while the backend also falls back to DIFFUSION_MODEL_BASE_MODELS. Models like Anima (file type "Model") were offered checkpoint roots in the UI even though use_default_paths would route them to the unet root. - Extract the two-tier decision into py/services/download_routing.py and reuse it in DownloadManager._execute_download - Add POST /api/lm/download/routing so the UI asks the backend for the routing decision; fall back to the local file-type check on failure - ModelVersionsTab: search both checkpoint and unet roots when resolving an existing version's download path
This commit is contained in:
@@ -20,7 +20,6 @@ from urllib.parse import urlparse
|
||||
from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
|
||||
from ..utils.constants import (
|
||||
CARD_PREVIEW_WIDTH,
|
||||
DIFFUSION_MODEL_BASE_MODELS,
|
||||
MODEL_WEIGHT_FILE_TYPES,
|
||||
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
|
||||
VALID_LORA_TYPES,
|
||||
@@ -32,6 +31,7 @@ from ..utils.utils import sanitize_folder_name
|
||||
from ..utils.exif_utils import ExifUtils
|
||||
from ..utils.metadata_manager import MetadataManager
|
||||
from .service_registry import ServiceRegistry
|
||||
from .download_routing import is_diffusion_model_download
|
||||
from .settings_manager import get_settings_manager
|
||||
from .metadata_service import get_default_metadata_provider, get_metadata_provider
|
||||
from .downloader import get_downloader, DownloadProgress, DownloadStreamControl
|
||||
@@ -1621,27 +1621,13 @@ class DownloadManager:
|
||||
}
|
||||
|
||||
# Check if this checkpoint should be treated as a diffusion model
|
||||
# Priority: (1) any file has type "UNet" or "Diffusion Model",
|
||||
# (2) baseModel is in DIFFUSION_MODEL_BASE_MODELS
|
||||
is_diffusion_model = False
|
||||
if model_type == "checkpoint":
|
||||
# Check file types first (more direct signal from CivitAI)
|
||||
version_files = version_info.get("files", [])
|
||||
for f in version_files:
|
||||
f_type = f.get("type", "")
|
||||
if f_type in ("UNet", "Diffusion Model"):
|
||||
is_diffusion_model = True
|
||||
logger.info(
|
||||
f"File type '{f_type}' detected, routing checkpoint to unet folder"
|
||||
)
|
||||
break
|
||||
|
||||
# Fallback to baseModel name check
|
||||
if not is_diffusion_model and 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"
|
||||
)
|
||||
# (shared with the download routing endpoint so the UI location
|
||||
# step and the actual download agree on the target roots).
|
||||
is_diffusion_model = is_diffusion_model_download(
|
||||
model_type,
|
||||
file_types=(f.get("type", "") for f in version_info.get("files", [])),
|
||||
base_model=base_model_value,
|
||||
)
|
||||
|
||||
# Existence check after the metadata fetch (#1058):
|
||||
# - An explicit file selection only blocks when THIS file is
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Shared download routing logic.
|
||||
|
||||
Decides whether a download initiated from the checkpoint library should be
|
||||
routed to the unet/diffusion-model roots instead of the checkpoint roots.
|
||||
Used by both the download manager (at download time) and the download
|
||||
routing HTTP endpoint (when the user picks a location in the UI), so the
|
||||
two can never disagree.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Iterable
|
||||
|
||||
from ..utils.constants import DIFFUSION_MODEL_BASE_MODELS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# File types reported by the CivitAI API that indicate a raw diffusion
|
||||
# model (loaded via UNETLoader in ComfyUI) rather than a full checkpoint.
|
||||
DIFFUSION_FILE_TYPES = frozenset({"UNet", "Diffusion Model"})
|
||||
|
||||
|
||||
def is_diffusion_model_download(
|
||||
model_type: str,
|
||||
file_types: Iterable[str] = (),
|
||||
base_model: str = "",
|
||||
) -> bool:
|
||||
"""Return True when a download should be routed to the unet roots.
|
||||
|
||||
Only applies to downloads initiated from the checkpoint library.
|
||||
Priority: (1) any file has type "UNet" or "Diffusion Model" (the more
|
||||
direct signal from CivitAI), (2) baseModel is a known diffusion model.
|
||||
"""
|
||||
if model_type != "checkpoint":
|
||||
return False
|
||||
|
||||
for file_type in file_types:
|
||||
if file_type in DIFFUSION_FILE_TYPES:
|
||||
logger.info(
|
||||
"File type '%s' detected, routing checkpoint to unet folder",
|
||||
file_type,
|
||||
)
|
||||
return True
|
||||
|
||||
if base_model in DIFFUSION_MODEL_BASE_MODELS:
|
||||
logger.info(
|
||||
"baseModel '%s' is a known diffusion model, routing to unet folder",
|
||||
base_model,
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
Reference in New Issue
Block a user