diff --git a/py/services/civitai_base_model_service.py b/py/services/civitai_base_model_service.py index e0dc89b9..4bce7780 100644 --- a/py/services/civitai_base_model_service.py +++ b/py/services/civitai_base_model_service.py @@ -196,6 +196,7 @@ class CivitaiBaseModelService: "ernie": "ERNI", "ernie turbo": "ETRB", "nucleus": "NUCL", + "krea 2": "KR2", "svd": "SVD", "ltxv": "LTXV", "ltxv2": "LTV2", @@ -424,6 +425,7 @@ class CivitaiBaseModelService: "Ernie", "Ernie Turbo", "Nucleus", + "Krea 2", ], } diff --git a/py/utils/constants.py b/py/utils/constants.py index 95d3ddbc..ec569fab 100644 --- a/py/utils/constants.py +++ b/py/utils/constants.py @@ -229,5 +229,6 @@ SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS = frozenset( "Ernie", "Ernie Turbo", "Nucleus", + "Krea 2", ] ) diff --git a/static/js/components/shared/ModelMetadata.js b/static/js/components/shared/ModelMetadata.js index f7c4e0ab..cada0bae 100644 --- a/static/js/components/shared/ModelMetadata.js +++ b/static/js/components/shared/ModelMetadata.js @@ -3,7 +3,7 @@ * Handles model metadata editing functionality - General version */ -import { BASE_MODEL_CATEGORIES } from '../../utils/constants.js'; +import { BASE_MODEL_CATEGORIES, getMergedBaseModels } from '../../utils/constants.js'; import { showToast } from '../../utils/uiHelpers.js'; import { getModelApiClient } from '../../api/modelApiFactory.js'; @@ -267,6 +267,7 @@ export function setupBaseModelEditing(filePath) { // Add options from BASE_MODEL_CATEGORIES constants const baseModelCategories = BASE_MODEL_CATEGORIES; + const categorizedModels = new Set(); // Create option groups for better organization Object.entries(baseModelCategories).forEach(([category, models]) => { @@ -277,13 +278,30 @@ export function setupBaseModelEditing(filePath) { const option = document.createElement('option'); option.value = model; option.textContent = model; - option.selected = model === currentValue; + if (model === currentValue) option.selected = true; + categorizedModels.add(model); group.appendChild(option); }); dropdown.appendChild(group); }); + // Check for dynamic base models from API that aren't in any category + const mergedModels = getMergedBaseModels(); + const uncategorizedModels = mergedModels.filter(model => !categorizedModels.has(model)); + if (uncategorizedModels.length > 0) { + const group = document.createElement('optgroup'); + group.label = 'Other (API)'; + uncategorizedModels.forEach(model => { + const option = document.createElement('option'); + option.value = model; + option.textContent = model; + if (model === currentValue) option.selected = true; + group.appendChild(option); + }); + dropdown.appendChild(group); + } + // Replace content with dropdown baseModelContent.style.display = 'none'; baseModelDisplay.insertBefore(dropdown, editBtn); diff --git a/static/js/utils/constants.js b/static/js/utils/constants.js index 3a6f3039..33da9d6e 100644 --- a/static/js/utils/constants.js +++ b/static/js/utils/constants.js @@ -70,6 +70,7 @@ export const BASE_MODELS = { ERNIE_TURBO: "Ernie Turbo", NUCLEUS: "Nucleus", PONY_V7: "Pony V7", + KREA_2: "Krea 2", // Default UNKNOWN: "Other" }; @@ -197,6 +198,7 @@ export const BASE_MODEL_ABBREVIATIONS = { [BASE_MODELS.ERNIE]: 'ERNI', [BASE_MODELS.ERNIE_TURBO]: 'ETRB', [BASE_MODELS.NUCLEUS]: 'NUCL', + [BASE_MODELS.KREA_2]: 'KR2', // Default [BASE_MODELS.UNKNOWN]: 'OTH' @@ -401,6 +403,7 @@ export const BASE_MODEL_CATEGORIES = { BASE_MODELS.PIXART_A, BASE_MODELS.PIXART_E, BASE_MODELS.HUNYUAN_1, BASE_MODELS.LUMINA, BASE_MODELS.KOLORS, BASE_MODELS.NOOBAI, BASE_MODELS.ANIMA, BASE_MODELS.ERNIE, BASE_MODELS.ERNIE_TURBO, BASE_MODELS.NUCLEUS, + BASE_MODELS.KREA_2, BASE_MODELS.UNKNOWN ] };