diff --git a/py/services/download_manager.py b/py/services/download_manager.py index adcde582..a38fb995 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -1288,10 +1288,24 @@ class DownloadManager: "download_id": download_id, } - # Check if this checkpoint should be treated as a diffusion model based on baseModel + # 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": - if base_model_value in DIFFUSION_MODEL_BASE_MODELS: + # 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" @@ -1420,7 +1434,7 @@ class DownloadManager: f for f in files if f.get("primary") - and f.get("type") in ("Model", "Negative", "Diffusion Model") + and f.get("type") in ("Model", "Negative", "Diffusion Model", "UNet") ), None, ) @@ -1451,7 +1465,7 @@ class DownloadManager: ( f for f in files - if f.get("primary") and f.get("type") in ("Model", "Negative", "Diffusion Model") + if f.get("primary") and f.get("type") in ("Model", "Negative", "Diffusion Model", "UNet") ), None, ) diff --git a/py/utils/constants.py b/py/utils/constants.py index cf812a06..b9e5f2ba 100644 --- a/py/utils/constants.py +++ b/py/utils/constants.py @@ -147,6 +147,8 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset( "Qwen", "ZImageBase", "ZImageTurbo", + # Krea 2 — loaded via UNETLoader in ComfyUI + "Krea 2", ] ) diff --git a/static/js/managers/DownloadManager.js b/static/js/managers/DownloadManager.js index 6c6e3e19..c12ddfab 100644 --- a/static/js/managers/DownloadManager.js +++ b/static/js/managers/DownloadManager.js @@ -351,7 +351,7 @@ export class DownloadManager { const thumbnailUrl = firstImage ? firstImage.url : '/loras_static/images/no-preview.png'; // Count model-type files per version - const modelFiles = (version.files || []).filter(f => f.type === 'Model'); + const modelFiles = (version.files || []).filter(f => f.type === 'Model' || f.type === 'UNet' || f.type === 'Diffusion Model'); const primaryFile = modelFiles.find(f => f.primary) || modelFiles[0] || {}; const fileSize = version.modelSizeKB ? (version.modelSizeKB / 1024).toFixed(2) : @@ -478,7 +478,7 @@ export class DownloadManager { if (!version) return; this.currentVersion = version; - const modelFiles = (version.files || []).filter(f => f.type === 'Model'); + const modelFiles = (version.files || []).filter(f => f.type === 'Model' || f.type === 'UNet' || f.type === 'Diffusion Model'); document.getElementById('versionStep').style.display = 'none'; document.getElementById('fileSelectionStep').style.display = 'block'; @@ -534,7 +534,7 @@ export class DownloadManager { const version = this.currentVersion; if (!version) return; - const modelFiles = (version.files || []).filter(f => f.type === 'Model'); + const modelFiles = (version.files || []).filter(f => f.type === 'Model' || f.type === 'UNet' || f.type === 'Diffusion Model'); this.selectedFile = modelFiles.find(f => f.id.toString() === selectedRadio.value); document.getElementById('fileSelectionStep').style.display = 'none'; @@ -954,7 +954,7 @@ export class DownloadManager { } if (!this.isBatchMode) { const fileParams = this.selectedFile ? { - type: 'Model', + type: this.selectedFile.type || 'Model', format: this.selectedFile.metadata?.format || 'SafeTensor', size: this.selectedFile.metadata?.size || 'full', fp: this.selectedFile.metadata?.fp,