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,