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:
Will Miao
2026-09-11 12:41:03 +08:00
parent 3cdc5ba7a2
commit e0052cd237
12 changed files with 479 additions and 30 deletions
@@ -0,0 +1,59 @@
"""HTTP handler for download target routing decisions."""
from __future__ import annotations
import json
import logging
from aiohttp import web
from ...services.download_routing import is_diffusion_model_download
logger = logging.getLogger(__name__)
class DownloadRoutingHandler:
"""Expose the download-time checkpoint/diffusion-model routing decision.
The web UI calls this when the user reaches the download location step
so the root dropdown offers the same root set (checkpoint vs unet) that
the download manager would pick for ``use_default_paths``.
"""
async def get_download_routing(self, request: web.Request) -> web.Response:
try:
payload = await request.json()
except json.JSONDecodeError:
return web.json_response(
{"success": False, "error": "Invalid JSON payload"}, status=400
)
model_type = payload.get("model_type", "")
base_model = payload.get("base_model") or ""
file_types = payload.get("file_types") or []
if not isinstance(model_type, str) or not model_type:
return web.json_response(
{"success": False, "error": "model_type is required"}, status=400
)
if not isinstance(base_model, str) or not isinstance(file_types, list):
return web.json_response(
{
"success": False,
"error": "base_model must be a string and file_types a list",
},
status=400,
)
is_diffusion = is_diffusion_model_download(
model_type,
file_types=(str(t) for t in file_types),
base_model=base_model,
)
return web.json_response(
{
"success": True,
"is_diffusion_model": is_diffusion,
"root_kind": "unet" if is_diffusion else model_type,
}
)
+5
View File
@@ -56,6 +56,7 @@ from ...utils.constants import (
)
from .hf_handlers import HfHandler
from .agent_handlers import AgentHandler
from .download_routing_handlers import DownloadRoutingHandler
from .model_handlers import ModelCivitaiHandler
from ...utils.civitai_utils import rewrite_preview_url
from ...utils.example_images_paths import (
@@ -3884,6 +3885,7 @@ class MiscHandlerSet:
base_model: BaseModelHandlerSet,
hf_handler: Any = None,
agent_handler: Any = None,
download_routing: Any = None,
) -> None:
self.health = health
self.settings = settings
@@ -3904,6 +3906,7 @@ class MiscHandlerSet:
self.base_model = base_model
self.hf_handler = hf_handler
self.agent_handler = agent_handler
self.download_routing = download_routing
def to_route_mapping(
self,
@@ -3962,6 +3965,8 @@ class MiscHandlerSet:
"get_agent_skills": self.agent_handler.get_agent_skills,
"execute_agent_skill": self.agent_handler.execute_agent_skill,
"cancel_agent_skill": self.agent_handler.cancel_agent_skill,
# Download routing handler
"get_download_routing": self.download_routing.get_download_routing,
# Base model handlers
"get_base_models": self.base_model.get_base_models,
"refresh_base_models": self.base_model.refresh_base_models,
+4
View File
@@ -103,6 +103,10 @@ MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition(
"GET", "/api/lm/hf-repo-files", "get_hf_repo_files"
),
# Download target routing decision (checkpoint vs diffusion model roots)
RouteDefinition(
"POST", "/api/lm/download/routing", "get_download_routing"
),
RouteDefinition(
"POST", "/api/lm/download-hf-model", "download_hf_model"
),
+3
View File
@@ -41,6 +41,7 @@ from .handlers.misc_handlers import (
from .handlers.base_model_handlers import BaseModelHandlerSet
from .handlers.hf_handlers import HfHandler
from .handlers.agent_handlers import AgentHandler
from .handlers.download_routing_handlers import DownloadRoutingHandler
from .misc_route_registrar import MiscRouteRegistrar
logger = logging.getLogger(__name__)
@@ -140,6 +141,7 @@ class MiscRoutes:
base_model = BaseModelHandlerSet()
hf_handler = HfHandler()
agent_handler = AgentHandler()
download_routing = DownloadRoutingHandler()
return self._handler_set_factory(
health=health,
@@ -161,6 +163,7 @@ class MiscRoutes:
base_model=base_model,
hf_handler=hf_handler,
agent_handler=agent_handler,
download_routing=download_routing,
)
+8 -22
View File
@@ -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
+53
View File
@@ -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
+1
View File
@@ -184,6 +184,7 @@ export const DOWNLOAD_ENDPOINTS = {
downloadGet: '/api/lm/download-model-get',
cancelGet: '/api/lm/cancel-download-get',
progress: '/api/lm/download-progress',
routing: '/api/lm/download/routing',
exampleImages: '/api/lm/force-download-example-images', // Re-process example images ignoring previous status
exampleImagesMissing: '/api/lm/download-example-images' // Download only missing example images
};
@@ -1390,8 +1390,22 @@ export function initVersionsTab({
try {
const client = ensureClient();
const rootsData = await client.fetchModelRoots();
const roots = rootsData?.roots;
// On the checkpoints page a diffusion model lives under the unet
// roots, so both root sets are needed to locate the current file.
let roots;
if (modelType === 'checkpoints') {
const [checkpointRoots, unetRoots] = await Promise.all([
client.fetchModelRoots(),
client.fetchModelRoots('diffusion_model'),
]);
roots = [
...(checkpointRoots?.roots || []),
...(unetRoots?.roots || []),
];
} else {
const rootsData = await client.fetchModelRoots();
roots = rootsData?.roots;
}
if (!Array.isArray(roots) || roots.length === 0) {
return null;
}
+51 -6
View File
@@ -3,6 +3,7 @@ import { showToast, setupAutoNewlineOnPaste } from '../utils/uiHelpers.js';
import { state } from '../state/index.js';
import { LoadingManager } from './LoadingManager.js';
import { getModelApiClient, resetAndReload } from '../api/modelApiFactory.js';
import { DOWNLOAD_ENDPOINTS } from '../api/apiConfig.js';
import { isModelWeightFile } from '../utils/modelFileTypes.js';
import { getStorageItem, setStorageItem } from '../utils/storageHelpers.js';
import { FolderTreeManager } from '../components/FolderTreeManager.js';
@@ -954,12 +955,7 @@ export class DownloadManager {
async proceedToLocationContent() {
try {
const _isDiffusionModel = this.selectedFile
? (this.selectedFile.type === 'UNet' || this.selectedFile.type === 'Diffusion Model')
: (this.currentVersion?.files || []).some(
f => f.type === 'UNet' || f.type === 'Diffusion Model'
);
this._isDiffusionModel = _isDiffusionModel;
this._isDiffusionModel = await this._resolveIsDiffusionModel();
let rootsData;
if (this._isDiffusionModel && this.apiClient.modelType === 'checkpoints') {
@@ -1020,6 +1016,55 @@ export class DownloadManager {
}
}
/**
* Decide whether this download routes to the diffusion model (unet)
* roots rather than the checkpoint roots. The backend owns the routing
* rule (file type first, baseModel fallback), so the location step asks
* it; if the endpoint is unavailable we degrade to the local file-type
* signal, which matches the backend for well-annotated models.
*/
async _resolveIsDiffusionModel() {
const localFileTypeCheck = this.selectedFile
? (this.selectedFile.type === 'UNet' || this.selectedFile.type === 'Diffusion Model')
: (this.currentVersion?.files || []).some(
f => f.type === 'UNet' || f.type === 'Diffusion Model'
);
// Only checkpoint downloads can route to the diffusion model roots;
// without version metadata (e.g. Hugging Face downloads) the local
// signal is all we have.
if (this.apiClient.modelType !== 'checkpoints'
|| (!this.selectedFile && !this.currentVersion)) {
return localFileTypeCheck;
}
try {
const fileTypes = this.selectedFile
? [this.selectedFile.type]
: (this.currentVersion?.files || []).map(f => f.type);
const response = await fetch(DOWNLOAD_ENDPOINTS.routing, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
model_type: 'checkpoint',
base_model: this.currentVersion?.baseModel || '',
file_types: fileTypes,
}),
});
if (!response.ok) {
throw new Error(`routing endpoint returned ${response.status}`);
}
const data = await response.json();
if (typeof data.is_diffusion_model === 'boolean') {
return data.is_diffusion_model;
}
} catch (error) {
console.warn('[download] routing endpoint unavailable, '
+ 'falling back to local file-type check:', error);
}
return localFileTypeCheck;
}
loadDefaultPathSetting() {
const modelType = this.apiClient.modelType;
const storageKey = `use_default_path_${modelType}`;
@@ -0,0 +1,146 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
const {
DOWNLOAD_MANAGER_MODULE,
MODAL_MANAGER_MODULE,
UI_HELPERS_MODULE,
STATE_MODULE,
LOADING_MANAGER_MODULE,
API_FACTORY_MODULE,
STORAGE_HELPERS_MODULE,
FOLDER_TREE_MANAGER_MODULE,
I18N_HELPERS_MODULE,
SUMMARY_MODULE,
} = vi.hoisted(() => ({
DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname,
MODAL_MANAGER_MODULE: new URL('../../../static/js/managers/ModalManager.js', import.meta.url).pathname,
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
LOADING_MANAGER_MODULE: new URL('../../../static/js/managers/LoadingManager.js', import.meta.url).pathname,
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
STORAGE_HELPERS_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
FOLDER_TREE_MANAGER_MODULE: new URL('../../../static/js/components/FolderTreeManager.js', import.meta.url).pathname,
I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
SUMMARY_MODULE: new URL('../../../static/js/components/DownloadBatchSummaryModal.js', import.meta.url).pathname,
}));
vi.mock(MODAL_MANAGER_MODULE, () => ({
modalManager: { showModal: vi.fn(), closeModal: vi.fn() },
}));
vi.mock(UI_HELPERS_MODULE, () => ({
showToast: vi.fn(),
setupAutoNewlineOnPaste: vi.fn(),
}));
vi.mock(STATE_MODULE, () => ({
state: { global: { settings: {} }, loadingManager: {} },
}));
vi.mock(LOADING_MANAGER_MODULE, () => ({
LoadingManager: vi.fn(() => ({})),
}));
vi.mock(API_FACTORY_MODULE, () => ({
getModelApiClient: vi.fn(),
resetAndReload: vi.fn(),
}));
vi.mock(STORAGE_HELPERS_MODULE, () => ({
getStorageItem: vi.fn((_key, defaultValue) => defaultValue),
setStorageItem: vi.fn(),
}));
vi.mock(FOLDER_TREE_MANAGER_MODULE, () => ({
FolderTreeManager: vi.fn(() => ({})),
}));
vi.mock(I18N_HELPERS_MODULE, () => ({
translate: vi.fn((_key, _vars, fallback) => fallback ?? ''),
}));
vi.mock(SUMMARY_MODULE, () => ({
showDownloadBatchSummary: vi.fn(),
}));
const { DownloadManager } = await import(DOWNLOAD_MANAGER_MODULE);
describe('DownloadManager._resolveIsDiffusionModel', () => {
let manager;
let fetchMock;
beforeEach(() => {
manager = new DownloadManager();
manager.apiClient = { modelType: 'checkpoints' };
manager.selectedFile = null;
manager.selectedFiles = [];
manager.currentVersion = null;
fetchMock = vi.fn();
vi.stubGlobal('fetch', fetchMock);
});
afterEach(() => {
vi.unstubAllGlobals();
});
function mockRoutingResponse(data, ok = true) {
fetchMock.mockResolvedValue({
ok,
status: ok ? 200 : 500,
json: async () => data,
});
}
it('asks the backend and routes baseModel-only diffusion models to unet roots', async () => {
// The reported Anima case: file type is plain "Model".
manager.currentVersion = { baseModel: 'Anima', files: [{ type: 'Model' }] };
mockRoutingResponse({ success: true, is_diffusion_model: true, root_kind: 'unet' });
expect(await manager._resolveIsDiffusionModel()).toBe(true);
expect(fetchMock).toHaveBeenCalledWith('/api/lm/download/routing', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
model_type: 'checkpoint',
base_model: 'Anima',
file_types: ['Model'],
}),
});
});
it('returns the backend decision for regular checkpoints', async () => {
manager.currentVersion = { baseModel: 'SDXL 1.0', files: [{ type: 'Model' }] };
mockRoutingResponse({ success: true, is_diffusion_model: false, root_kind: 'checkpoint' });
expect(await manager._resolveIsDiffusionModel()).toBe(false);
});
it('sends only the selected file type when a file is selected', async () => {
manager.currentVersion = { baseModel: 'Flux.1 D', files: [{ type: 'Model' }, { type: 'UNet' }] };
manager.selectedFile = { type: 'UNet' };
mockRoutingResponse({ success: true, is_diffusion_model: true, root_kind: 'unet' });
expect(await manager._resolveIsDiffusionModel()).toBe(true);
expect(JSON.parse(fetchMock.mock.calls[0][1].body).file_types).toEqual(['UNet']);
});
it('falls back to the local file-type check when the endpoint fails', async () => {
manager.currentVersion = { baseModel: 'SDXL 1.0', files: [{ type: 'UNet' }] };
fetchMock.mockRejectedValue(new Error('network down'));
expect(await manager._resolveIsDiffusionModel()).toBe(true);
});
it('falls back to false when the endpoint fails and no local signal exists', async () => {
manager.currentVersion = { baseModel: 'Anima', files: [{ type: 'Model' }] };
mockRoutingResponse({}, false);
expect(await manager._resolveIsDiffusionModel()).toBe(false);
});
it('never calls the endpoint for non-checkpoint pages', async () => {
manager.apiClient = { modelType: 'loras' };
manager.currentVersion = { baseModel: 'Anima', files: [{ type: 'Model' }] };
expect(await manager._resolveIsDiffusionModel()).toBe(false);
expect(fetchMock).not.toHaveBeenCalled();
});
it('never calls the endpoint without version metadata (e.g. Hugging Face)', async () => {
expect(await manager._resolveIsDiffusionModel()).toBe(false);
expect(fetchMock).not.toHaveBeenCalled();
});
});
@@ -0,0 +1,93 @@
"""Tests for the download routing HTTP handler."""
import json
import pytest
from py.routes.handlers.download_routing_handlers import DownloadRoutingHandler
class FakeRequest:
def __init__(self, payload):
self._payload = payload
async def json(self):
if isinstance(self._payload, Exception):
raise self._payload
return self._payload
@pytest.mark.asyncio
async def test_diffusion_base_model_routes_to_unet():
"""The reported Anima case: file type "Model", baseModel "Anima"."""
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(
FakeRequest(
{"model_type": "checkpoint", "base_model": "Anima", "file_types": ["Model"]}
)
)
payload = json.loads(response.text)
assert response.status == 200
assert payload == {"success": True, "is_diffusion_model": True, "root_kind": "unet"}
@pytest.mark.asyncio
async def test_unet_file_type_routes_to_unet():
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(
FakeRequest(
{"model_type": "checkpoint", "base_model": "SDXL 1.0", "file_types": ["UNet"]}
)
)
payload = json.loads(response.text)
assert payload["is_diffusion_model"] is True
assert payload["root_kind"] == "unet"
@pytest.mark.asyncio
async def test_regular_checkpoint_stays_on_checkpoint_root():
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(
FakeRequest(
{"model_type": "checkpoint", "base_model": "SDXL 1.0", "file_types": ["Model"]}
)
)
payload = json.loads(response.text)
assert payload["is_diffusion_model"] is False
assert payload["root_kind"] == "checkpoint"
@pytest.mark.asyncio
async def test_lora_is_never_diffusion():
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(
FakeRequest({"model_type": "lora", "base_model": "Anima", "file_types": []})
)
payload = json.loads(response.text)
assert payload["is_diffusion_model"] is False
assert payload["root_kind"] == "lora"
@pytest.mark.asyncio
async def test_missing_model_type_rejected():
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(FakeRequest({"base_model": "Anima"}))
assert response.status == 400
@pytest.mark.asyncio
async def test_invalid_file_types_rejected():
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(
FakeRequest({"model_type": "checkpoint", "file_types": "Model"})
)
assert response.status == 400
@pytest.mark.asyncio
async def test_invalid_json_rejected():
handler = DownloadRoutingHandler()
response = await handler.get_download_routing(
FakeRequest(json.JSONDecodeError("bad", "", 0))
)
assert response.status == 400
+40
View File
@@ -0,0 +1,40 @@
"""Tests for the shared download routing decision."""
import pytest
from py.services.download_routing import is_diffusion_model_download
@pytest.mark.parametrize("file_type", ["UNet", "Diffusion Model"])
def test_file_type_signal_routes_to_unet(file_type):
assert is_diffusion_model_download(
"checkpoint", file_types=[file_type], base_model="SDXL 1.0"
)
def test_base_model_fallback_routes_to_unet():
"""The reported Anima case: file type is plain "Model", but the
baseModel is a known diffusion model."""
assert is_diffusion_model_download(
"checkpoint", file_types=["Model"], base_model="Anima"
)
def test_regular_checkpoint_stays_on_checkpoint_roots():
assert not is_diffusion_model_download(
"checkpoint", file_types=["Model"], base_model="SDXL 1.0"
)
def test_non_checkpoint_types_never_route_to_unet():
assert not is_diffusion_model_download(
"lora", file_types=["UNet"], base_model="Anima"
)
assert not is_diffusion_model_download(
"embedding", file_types=["Diffusion Model"], base_model="Anima"
)
def test_empty_inputs_stay_on_checkpoint_roots():
assert not is_diffusion_model_download("checkpoint")
assert not is_diffusion_model_download("checkpoint", file_types=[], base_model="")