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,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