mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 03:01:27 -03:00
Will Miao
e0052cd237
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
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
"""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
|