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
+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}`;