fix(download): accept UNet and Diffusion Model file types from CivitAI

- Prefer file type (UNet/Diffusion Model) over baseModel name when
  deciding whether a checkpoint routes to the unet folder
- Add UNet to backend primary file type whitelist
- Add Krea 2 to DIFFUSION_MODEL_BASE_MODELS
- Include UNet/Diffusion Model files in frontend file selection UI
- Use actual file type from CivitAI in download params instead of
  hardcoded 'Model'

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

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",
     ]
 )
 

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,