mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 03:01:27 -03:00
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:
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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}`;
|
||||
|
||||
Reference in New Issue
Block a user