From e0052cd237ef2d67a5197e641414971843f2b8ae Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 11 Sep 2026 12:41:03 +0800 Subject: [PATCH] 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 --- .../handlers/download_routing_handlers.py | 59 +++++++ py/routes/handlers/misc_handlers.py | 5 + py/routes/misc_route_registrar.py | 4 + py/routes/misc_routes.py | 3 + py/services/download_manager.py | 30 +--- py/services/download_routing.py | 53 +++++++ static/js/api/apiConfig.js | 1 + .../js/components/shared/ModelVersionsTab.js | 18 ++- static/js/managers/DownloadManager.js | 57 ++++++- .../managers/downloadManager.routing.test.js | 146 ++++++++++++++++++ .../routes/test_download_routing_handlers.py | 93 +++++++++++ tests/services/test_download_routing.py | 40 +++++ 12 files changed, 479 insertions(+), 30 deletions(-) create mode 100644 py/routes/handlers/download_routing_handlers.py create mode 100644 py/services/download_routing.py create mode 100644 tests/frontend/managers/downloadManager.routing.test.js create mode 100644 tests/routes/test_download_routing_handlers.py create mode 100644 tests/services/test_download_routing.py diff --git a/py/routes/handlers/download_routing_handlers.py b/py/routes/handlers/download_routing_handlers.py new file mode 100644 index 00000000..0d3c7a12 --- /dev/null +++ b/py/routes/handlers/download_routing_handlers.py @@ -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, + } + ) diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index 3b13e685..e879327b 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -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, diff --git a/py/routes/misc_route_registrar.py b/py/routes/misc_route_registrar.py index 87313a6e..1e7a1aeb 100644 --- a/py/routes/misc_route_registrar.py +++ b/py/routes/misc_route_registrar.py @@ -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" ), diff --git a/py/routes/misc_routes.py b/py/routes/misc_routes.py index f572a731..58cb8beb 100644 --- a/py/routes/misc_routes.py +++ b/py/routes/misc_routes.py @@ -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, ) diff --git a/py/services/download_manager.py b/py/services/download_manager.py index d28be353..2f52f337 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -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 diff --git a/py/services/download_routing.py b/py/services/download_routing.py new file mode 100644 index 00000000..ec09d24d --- /dev/null +++ b/py/services/download_routing.py @@ -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 diff --git a/static/js/api/apiConfig.js b/static/js/api/apiConfig.js index c577cbe2..dbecbab5 100644 --- a/static/js/api/apiConfig.js +++ b/static/js/api/apiConfig.js @@ -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 }; diff --git a/static/js/components/shared/ModelVersionsTab.js b/static/js/components/shared/ModelVersionsTab.js index 26d31e9f..5aff6fea 100644 --- a/static/js/components/shared/ModelVersionsTab.js +++ b/static/js/components/shared/ModelVersionsTab.js @@ -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; } diff --git a/static/js/managers/DownloadManager.js b/static/js/managers/DownloadManager.js index c5ab53fc..6831ee64 100644 --- a/static/js/managers/DownloadManager.js +++ b/static/js/managers/DownloadManager.js @@ -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}`; diff --git a/tests/frontend/managers/downloadManager.routing.test.js b/tests/frontend/managers/downloadManager.routing.test.js new file mode 100644 index 00000000..c4e7ab59 --- /dev/null +++ b/tests/frontend/managers/downloadManager.routing.test.js @@ -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(); + }); +}); diff --git a/tests/routes/test_download_routing_handlers.py b/tests/routes/test_download_routing_handlers.py new file mode 100644 index 00000000..48bcf767 --- /dev/null +++ b/tests/routes/test_download_routing_handlers.py @@ -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 diff --git a/tests/services/test_download_routing.py b/tests/services/test_download_routing.py new file mode 100644 index 00000000..2d9dfb22 --- /dev/null +++ b/tests/services/test_download_routing.py @@ -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="")