diff --git a/locales/de.json b/locales/de.json index 37761e48..ebbc260d 100644 --- a/locales/de.json +++ b/locales/de.json @@ -233,6 +233,7 @@ "noCreditRequired": "Kein Credit erforderlich", "allowSellingGeneratedContent": "Verkauf erlaubt", "noTags": "Keine Tags", + "autoTags": "Auto-Tags", "noBaseModelMatches": "Keine Basismodelle entsprechen der aktuellen Suche.", "clearAll": "Alle Filter löschen", "any": "Beliebig", diff --git a/locales/en.json b/locales/en.json index 578fe089..b025b9e5 100644 --- a/locales/en.json +++ b/locales/en.json @@ -233,6 +233,7 @@ "noCreditRequired": "No Credit Required", "allowSellingGeneratedContent": "Allow Selling", "noTags": "No tags", + "autoTags": "Auto Tags", "noBaseModelMatches": "No base models match the current search.", "clearAll": "Clear All Filters", "any": "Any", diff --git a/locales/es.json b/locales/es.json index 38194a4f..36825d67 100644 --- a/locales/es.json +++ b/locales/es.json @@ -233,6 +233,7 @@ "noCreditRequired": "Sin crédito requerido", "allowSellingGeneratedContent": "Venta permitida", "noTags": "Sin etiquetas", + "autoTags": "Etiquetas automáticas", "noBaseModelMatches": "Ningún modelo base coincide con la búsqueda actual.", "clearAll": "Limpiar todos los filtros", "any": "Cualquiera", diff --git a/locales/fr.json b/locales/fr.json index 38ef9c57..7a02ffd3 100644 --- a/locales/fr.json +++ b/locales/fr.json @@ -233,6 +233,7 @@ "noCreditRequired": "Crédit non requis", "allowSellingGeneratedContent": "Vente autorisée", "noTags": "Aucun tag", + "autoTags": "Auto-Tags", "noBaseModelMatches": "Aucun modèle de base ne correspond à la recherche actuelle.", "clearAll": "Effacer tous les filtres", "any": "N'importe quel", diff --git a/locales/he.json b/locales/he.json index 43e5ba17..8e7c2a1f 100644 --- a/locales/he.json +++ b/locales/he.json @@ -233,6 +233,7 @@ "noCreditRequired": "ללא קרדיט נדרש", "allowSellingGeneratedContent": "אפשר מכירה", "noTags": "ללא תגיות", + "autoTags": "תגיות אוטומטיות", "noBaseModelMatches": "אין מודלי בסיס התואמים לחיפוש הנוכחי.", "clearAll": "נקה את כל המסננים", "any": "כלשהו", diff --git a/locales/ja.json b/locales/ja.json index e156d23f..45113b29 100644 --- a/locales/ja.json +++ b/locales/ja.json @@ -233,6 +233,7 @@ "noCreditRequired": "クレジット不要", "allowSellingGeneratedContent": "販売許可", "noTags": "タグなし", + "autoTags": "自動タグ", "noBaseModelMatches": "現在の検索に一致するベースモデルはありません。", "clearAll": "すべてのフィルタをクリア", "any": "いずれか", diff --git a/locales/ko.json b/locales/ko.json index 1fffcc68..ef4db6da 100644 --- a/locales/ko.json +++ b/locales/ko.json @@ -233,6 +233,7 @@ "noCreditRequired": "크레딧 표기 없음", "allowSellingGeneratedContent": "판매 허용", "noTags": "태그 없음", + "autoTags": "자동 태그", "noBaseModelMatches": "현재 검색과 일치하는 베이스 모델이 없습니다.", "clearAll": "모든 필터 지우기", "any": "아무", diff --git a/locales/ru.json b/locales/ru.json index 8a6f9ebe..dfbd4b08 100644 --- a/locales/ru.json +++ b/locales/ru.json @@ -233,6 +233,7 @@ "noCreditRequired": "Без указания авторства", "allowSellingGeneratedContent": "Продажа разрешена", "noTags": "Без тегов", + "autoTags": "Авто-теги", "noBaseModelMatches": "Нет базовых моделей, соответствующих текущему поиску.", "clearAll": "Очистить все фильтры", "any": "Любой", diff --git a/locales/zh-CN.json b/locales/zh-CN.json index 19db9752..e60aadb9 100644 --- a/locales/zh-CN.json +++ b/locales/zh-CN.json @@ -233,6 +233,7 @@ "noCreditRequired": "无需署名", "allowSellingGeneratedContent": "允许销售", "noTags": "无标签", + "autoTags": "自动标签", "noBaseModelMatches": "没有基础模型符合当前搜索。", "clearAll": "清除所有筛选", "any": "任一", diff --git a/locales/zh-TW.json b/locales/zh-TW.json index cf8c9a8f..cbe87428 100644 --- a/locales/zh-TW.json +++ b/locales/zh-TW.json @@ -233,6 +233,7 @@ "noCreditRequired": "無需署名", "allowSellingGeneratedContent": "允許銷售", "noTags": "無標籤", + "autoTags": "自動標籤", "noBaseModelMatches": "沒有基礎模型符合目前的搜尋。", "clearAll": "清除所有篩選", "any": "任一", diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py index 1748fe21..6e931b2d 100644 --- a/py/routes/handlers/model_handlers.py +++ b/py/routes/handlers/model_handlers.py @@ -301,6 +301,15 @@ class ModelListingHandler: for tag in exclude_tags: if tag: tag_filters[tag] = "exclude" + + auto_tag_filters: Dict[str, str] = {} + for tag in request.query.getall("auto_tag_include", []): + if tag: + auto_tag_filters[tag] = "include" + for tag in request.query.getall("auto_tag_exclude", []): + if tag: + auto_tag_filters[tag] = "exclude" + favorites_only = request.query.get("favorites_only", "false").lower() == "true" search_options = { @@ -367,6 +376,7 @@ class ModelListingHandler: "fuzzy_search": fuzzy_search, "base_models": base_models, "tags": tag_filters, + "auto_tags": auto_tag_filters, "tag_logic": tag_logic, "search_options": search_options, "hash_filters": hash_filters, diff --git a/py/services/auto_tag_service.py b/py/services/auto_tag_service.py new file mode 100644 index 00000000..545cf52e --- /dev/null +++ b/py/services/auto_tag_service.py @@ -0,0 +1,121 @@ +""" +Auto-tag extraction service for model cards. + +Extracts implicit model attributes (HIGH/LOW, I2V/T2V/TI2V, Lightning, Turbo) +from filename, base_model, and CivitAI version name — no manual tagging required. +""" + +from __future__ import annotations + +import re +from typing import Dict, List, Set + +# ── Tag category definitions ────────────────────────────────────────── +# Each category maps a display label to a regex pattern. +# Patterns are case-insensitive and matched against filename, base_model, +# and civitai version name. + +# Use (? List[str]: + """Collect all text sources from model data for tag matching.""" + sources: List[str] = [] + + file_name = model_data.get("file_name", "") + if file_name: + sources.append(file_name) + + base_model = model_data.get("base_model", "") + if base_model: + sources.append(base_model) + + civitai = model_data.get("civitai", {}) + if isinstance(civitai, dict): + version_name = civitai.get("name", "") + if version_name: + sources.append(version_name) + + return sources + + +def extract_auto_tags(model_data: Dict) -> List[str]: + """Extract auto-detected tags from model metadata. + + Matches predefined patterns against filename, base_model, and + CivitAI version name. Returns a sorted, deduplicated list of tag labels. + + HIGH/LOW tags are only returned when the base_model indicates a Wan + family model — no other model architecture uses this distinction. + + Args: + model_data: Model metadata dict with keys: + file_name, base_model, civitai (with optional 'name' field). + + Returns: + Sorted list of unique auto-tag strings (e.g. ["I2V"]). + """ + sources = _collect_sources(model_data) + if not sources: + return [] + + base_model = model_data.get("base_model", "") + is_wan = "wan" in base_model.lower() + + found: Set[str] = set() + + for label, pattern in AUTO_TAG_CATEGORIES.items(): + # HIGH/LOW are Wan-specific — skip for non-Wan to avoid noise + if label in ("HIGH", "LOW"): + if not is_wan: + continue + # Use case-insensitive character class + case-sensitive boundary, + # so "HighNoise" (camelCase) matches but "highlight" doesn't. + # Boundary: not followed by lowercase letter (= word has ended). + ci = "".join(f"[{c.lower()}{c.upper()}]" for c in label) + if label == "LOW": + regex = re.compile(r"(? Dict: diff --git a/py/services/embedding_service.py b/py/services/embedding_service.py index 71ad259a..ea187f85 100644 --- a/py/services/embedding_service.py +++ b/py/services/embedding_service.py @@ -3,6 +3,7 @@ import logging from typing import Dict from .base_model_service import BaseModelService +from .auto_tag_service import extract_auto_tags from ..utils.models import EmbeddingMetadata from ..config import config @@ -45,7 +46,8 @@ class EmbeddingService(BaseModelService): "exclude": bool(embedding_data.get("exclude", False)), "update_available": bool(embedding_data.get("update_available", False)), "skip_metadata_refresh": bool(embedding_data.get("skip_metadata_refresh", False)), - "civitai": self.filter_civitai_data(embedding_data.get("civitai", {}), minimal=True) + "civitai": self.filter_civitai_data(embedding_data.get("civitai", {}), minimal=True), + "auto_tags": embedding_data.get("auto_tags") or extract_auto_tags(embedding_data), } def find_duplicate_hashes(self) -> Dict: diff --git a/py/services/lora_service.py b/py/services/lora_service.py index d0d6b769..f1a121ac 100644 --- a/py/services/lora_service.py +++ b/py/services/lora_service.py @@ -5,6 +5,7 @@ from typing import Dict, List, Optional from .base_model_service import BaseModelService from .model_query import resolve_sub_type +from .auto_tag_service import extract_auto_tags from ..utils.models import LoraMetadata from ..config import config @@ -57,6 +58,7 @@ class LoraService(BaseModelService): "civitai": self.filter_civitai_data( lora_data.get("civitai", {}), minimal=True ), + "auto_tags": lora_data.get("auto_tags") or extract_auto_tags(lora_data), } async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]: diff --git a/py/services/model_query.py b/py/services/model_query.py index 0cc91c40..34c542d4 100644 --- a/py/services/model_query.py +++ b/py/services/model_query.py @@ -96,6 +96,7 @@ class FilterCriteria: folder_exclude: Optional[Sequence[str]] = None base_models: Optional[Sequence[str]] = None tags: Optional[Dict[str, str]] = None + auto_tags: Optional[Dict[str, str]] = None favorites_only: bool = False search_options: Optional[Dict[str, Any]] = None model_types: Optional[Sequence[str]] = None @@ -359,10 +360,37 @@ class ModelFilterSet: ] model_types_duration = time.perf_counter() - t0 + auto_tags_duration = 0 + auto_tag_filters = criteria.auto_tags or {} + if auto_tag_filters: + t0 = time.perf_counter() + include_at = set() + exclude_at = set() + for tag, state in auto_tag_filters.items(): + if not tag: + continue + if state == "exclude": + exclude_at.add(tag) + else: + include_at.add(tag) + + if include_at: + items = [ + item for item in items + if any(tag in include_at for tag in (item.get("auto_tags") or [])) + ] + + if exclude_at: + items = [ + item for item in items + if not any(tag in exclude_at for tag in (item.get("auto_tags") or [])) + ] + auto_tags_duration = time.perf_counter() - t0 + duration = time.perf_counter() - overall_start if duration > 0.1: # Only log if it's potentially slow logger.debug( - "ModelFilterSet.apply took %.3fs (sfw: %.3fs, fav: %.3fs, folder: %.3fs, base: %.3fs, tags: %.3fs, types: %.3fs). " + "ModelFilterSet.apply took %.3fs (sfw: %.3fs, fav: %.3fs, folder: %.3fs, base: %.3fs, tags: %.3fs, types: %.3fs, auto_tags: %.3fs). " "Count: %d -> %d", duration, sfw_duration, @@ -371,6 +399,7 @@ class ModelFilterSet: base_models_duration, tags_duration, model_types_duration, + auto_tags_duration, initial_count, len(items), ) diff --git a/static/css/components/card.css b/static/css/components/card.css index 5db655fe..cda903a7 100644 --- a/static/css/components/card.css +++ b/static/css/components/card.css @@ -507,21 +507,96 @@ background: rgba(0,0,0,0.18); /* Optional: subtle background for contrast */ } +/* Version row — flex container for badges + version names */ +.version-row { + display: flex; + flex-wrap: wrap; + align-items: center; + gap: 3px; + margin-top: 2px; +} + +/* Badge + version-name binding: they wrap as a single unit */ +.badge-version-unit { + display: inline-flex; + align-items: center; + gap: 3px; + min-width: 0; + flex-shrink: 0; +} + /* Medium density adjustments for version name */ .medium-density .version-name { font-size: 0.8em; } +.medium-density .badge-version-unit .version-name { + max-width: 90px; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} + /* Compact density adjustments for version name */ .compact-density .version-name { font-size: 0.75em; } -/* Hide civitai version name when setting is disabled */ -body.hide-card-version .civitai-version { +.compact-density .badge-version-unit .version-name { + max-width: 70px; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; +} + +.medium-density .version-row { + gap: 2px; +} + +/* HIGH / LOW badges — shown inline before version name in card footer */ +.hl-badge { + display: inline-block; + font-size: 0.7em; + font-weight: 600; + line-height: 1.1; + padding: 1px 5px; + border-radius: var(--border-radius-xs); + border: 1px solid rgba(255, 255, 255, 0.2); + white-space: nowrap; +} + +.hl-badge--high { + color: oklch(75% 0.12 230); + background: oklch(55% 0.15 240 / 0.25); + border-color: oklch(60% 0.18 250 / 0.3); +} + +.hl-badge--low { + color: oklch(78% 0.10 185); + background: oklch(50% 0.10 190 / 0.25); + border-color: oklch(55% 0.12 195 / 0.3); +} + +.medium-density .hl-badge { + font-size: 0.65em; +} + +.compact-density .hl-badge { + font-size: 0.62em; + padding: 0px 4px; +} + +/* Hide version-related elements when setting is disabled */ +body.hide-card-version .civitai-version, +body.hide-card-version .hl-badge { display: none; } +/* Compact density adjustments for version name */ +.compact-density .version-name { + font-size: 0.75em; +} + /* Prevent text selection on cards and interactive elements */ .model-card, .model-card *, diff --git a/static/js/api/baseModelApi.js b/static/js/api/baseModelApi.js index c22c8b39..bc5ec074 100644 --- a/static/js/api/baseModelApi.js +++ b/static/js/api/baseModelApi.js @@ -978,6 +978,16 @@ export class BaseModelApiClient { }); } + if (pageState.filters.autoTags && Object.keys(pageState.filters.autoTags).length > 0) { + Object.entries(pageState.filters.autoTags).forEach(([tag, state]) => { + if (state === 'include') { + params.append('auto_tag_include', tag); + } else if (state === 'exclude') { + params.append('auto_tag_exclude', tag); + } + }); + } + if (pageState.filters.baseModel && pageState.filters.baseModel.length > 0) { // Check for empty wildcard marker - if present, no models should match const EMPTY_WILDCARD_MARKER = '__EMPTY_WILDCARD_RESULT__'; diff --git a/static/js/components/shared/ModelCard.js b/static/js/components/shared/ModelCard.js index 4af2fb93..f9cc7e47 100644 --- a/static/js/components/shared/ModelCard.js +++ b/static/js/components/shared/ModelCard.js @@ -644,8 +644,23 @@ export function createModelCard(model, modelType) {
${getDisplayName(model)} -
- ${model.civitai?.name ? `${model.civitai.name}` : ''} +
+ ${(() => { + const autoTags = model.auto_tags || []; + const hlTags = autoTags.filter(t => t === 'HIGH' || t === 'LOW'); + const hasVersionName = model.civitai?.name; + if (!hlTags.length && !hasVersionName) return ''; + const density = state.global.settings.display_density || 'default'; + const shortLabels = density === 'medium' || density === 'compact'; + const badges = hlTags.map(t => { + const cls = t === 'HIGH' ? 'hl-badge hl-badge--high' : 'hl-badge hl-badge--low'; + const label = shortLabels ? (t === 'HIGH' ? 'H' : 'L') : t; + const titleAttr = shortLabels ? ` title="${t}"` : ''; + return `${label}`; + }).join(''); + const versionHtml = hasVersionName ? `${model.civitai.name}` : ''; + return `${badges}${versionHtml}`; + })()} ${hasUsageCount ? `${model.usage_count}×` : ''}
diff --git a/static/js/managers/FilterManager.js b/static/js/managers/FilterManager.js index 9af4d58e..d3de75a6 100644 --- a/static/js/managers/FilterManager.js +++ b/static/js/managers/FilterManager.js @@ -70,6 +70,9 @@ export class FilterManager { // Initialize tag logic toggle this.initializeTagLogicToggle(); + // Create auto-tag filter section (I2V, T2V, TI2V, Lightning, Turbo) + this.createAutoTagFilters(); + // Add click handler for filter button if (this.filterButton) { this.filterButton.addEventListener('click', () => { @@ -480,6 +483,58 @@ export class FilterManager { } } + AUTO_TAG_FILTER_TAGS = ['I2V', 'T2V', 'TI2V', 'Lightning', 'Turbo']; + + createAutoTagFilters() { + const container = document.getElementById('autoTagFilterTags'); + if (container) return; + + const modelTypeSection = document.getElementById('modelTypeTags')?.closest('.filter-section'); + if (!modelTypeSection) return; + + const section = document.createElement('div'); + section.className = 'filter-section'; + section.innerHTML = ` +

${translate('header.filter.autoTags', {}, 'Auto Tags')}

+
+ `; + modelTypeSection.parentNode.insertBefore(section, modelTypeSection.nextSibling); + + const tagsContainer = document.getElementById('autoTagFilterTags'); + this.AUTO_TAG_FILTER_TAGS.forEach(tag => { + const el = document.createElement('div'); + el.className = 'filter-tag auto-tag-filter'; + el.dataset.autoTag = tag; + el.textContent = tag; + + // Restore previous state + const state = (this.filters.autoTags && this.filters.autoTags[tag]) || 'none'; + this._applyTriState(el, state); + + el.addEventListener('click', async () => { + const current = (this.filters.autoTags && this.filters.autoTags[tag]) || 'none'; + const next = current === 'none' ? 'include' : current === 'include' ? 'exclude' : 'none'; + if (!this.filters.autoTags) this.filters.autoTags = {}; + if (next === 'none') { + delete this.filters.autoTags[tag]; + } else { + this.filters.autoTags[tag] = next; + } + this._applyTriState(el, next); + this.updateActiveFiltersCount(); + await this.applyFilters(false); + }); + + tagsContainer.appendChild(el); + }); + } + + _applyTriState(el, state) { + el.classList.remove('active', 'exclude'); + if (state === 'include') el.classList.add('active'); + else if (state === 'exclude') el.classList.add('exclude'); + } + toggleFilterPanel() { if (this.filterPanel) { const isHidden = this.filterPanel.classList.contains('hidden'); @@ -540,6 +595,13 @@ export class FilterManager { this.updateLicenseSelections(); } this.updateModelTypeSelections(); + + const autoTagEls = document.querySelectorAll('.auto-tag-filter'); + autoTagEls.forEach(el => { + const tag = el.dataset.autoTag; + const state = (this.filters.autoTags && this.filters.autoTags[tag]) || 'none'; + this._applyTriState(el, state); + }); } updateModelTypeSelections() { @@ -556,11 +618,12 @@ export class FilterManager { updateActiveFiltersCount() { const tagFilterCount = this.filters.tags ? Object.keys(this.filters.tags).length : 0; + const autoTagFilterCount = this.filters.autoTags ? Object.keys(this.filters.autoTags).length : 0; const licenseFilterCount = this.filters.license ? Object.keys(this.filters.license).length : 0; const modelTypeFilterCount = this.filters.modelTypes.length; // Exclude EMPTY_WILDCARD_MARKER from base model count const baseModelCount = this.filters.baseModel.filter(m => m !== EMPTY_WILDCARD_MARKER).length; - const totalActiveFilters = baseModelCount + tagFilterCount + licenseFilterCount + modelTypeFilterCount; + const totalActiveFilters = baseModelCount + tagFilterCount + autoTagFilterCount + licenseFilterCount + modelTypeFilterCount; if (this.activeFiltersCount) { if (totalActiveFilters > 0) { @@ -652,6 +715,7 @@ export class FilterManager { ...this.filters, baseModel: [], tags: {}, + autoTags: {}, license: {}, modelTypes: [], tagLogic: 'any' @@ -721,6 +785,7 @@ export class FilterManager { hasActiveFilters() { const tagCount = this.filters.tags ? Object.keys(this.filters.tags).length : 0; + const autoTagCount = this.filters.autoTags ? Object.keys(this.filters.autoTags).length : 0; const licenseCount = this.filters.license ? Object.keys(this.filters.license).length : 0; const modelTypeCount = this.filters.modelTypes.length; // Exclude EMPTY_WILDCARD_MARKER from base model count @@ -728,6 +793,7 @@ export class FilterManager { return ( baseModelCount > 0 || tagCount > 0 || + autoTagCount > 0 || licenseCount > 0 || modelTypeCount > 0 ); @@ -739,6 +805,7 @@ export class FilterManager { ...source, baseModel: Array.isArray(source.baseModel) ? [...source.baseModel] : [], tags: this.normalizeTagFilters(source.tags), + autoTags: this.normalizeTagFilters(source.autoTags), license: this.shouldShowLicenseFilters() ? this.normalizeLicenseFilters(source.license) : {}, modelTypes: this.normalizeModelTypeFilters(source.modelTypes), tagLogic: source.tagLogic || 'any' @@ -822,6 +889,7 @@ export class FilterManager { ...this.filters, baseModel: [...(this.filters.baseModel || [])], tags: { ...(this.filters.tags || {}) }, + autoTags: { ...(this.filters.autoTags || {}) }, license: { ...(this.filters.license || {}) }, modelTypes: [...(this.filters.modelTypes || [])], tagLogic: this.filters.tagLogic || 'any' diff --git a/static/js/utils/constants.js b/static/js/utils/constants.js index ceede86c..3a6f3039 100644 --- a/static/js/utils/constants.js +++ b/static/js/utils/constants.js @@ -500,6 +500,18 @@ export function clearDynamicBaseModels() { dynamicBaseModelsTimestamp = null; } +export const AUTO_TAG_GROUPS = { + mode: new Set(['HIGH', 'LOW']), + video: new Set(['I2V', 'T2V', 'TI2V']), + speed: new Set(['Lightning', 'Turbo']), +}; + +export const AUTO_TAG_GROUP_LABELS = { + mode: 'High / Low', + video: 'I2V / T2V / TI2V', + speed: 'Lightning / Turbo', +}; + /** * Check if dynamic base models cache is valid * @returns {boolean} diff --git a/tests/test_auto_tag_service.py b/tests/test_auto_tag_service.py new file mode 100644 index 00000000..e7c3fb45 --- /dev/null +++ b/tests/test_auto_tag_service.py @@ -0,0 +1,151 @@ +import pytest +import sys +import os + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "py")) + +from services.auto_tag_service import extract_auto_tags, AUTO_TAG_CATEGORIES + + +class TestExtractAutoTags: + def test_file_name_high_i2v(self): + result = extract_auto_tags({ + "file_name": "Shirt_lift_Wan2.2_14B_I2V_HIGH_v1.0", + "base_model": "Wan Video 2.2 I2V-A14B", + "civitai": {}, + }) + assert set(result) == {"HIGH", "I2V"} + + def test_file_name_t2v_low(self): + result = extract_auto_tags({ + "file_name": "my_wan_t2v_low_v2", + "base_model": "Wan 2.1", + "civitai": {}, + }) + assert set(result) == {"LOW", "T2V"} + + def test_file_name_ti2v_high(self): + result = extract_auto_tags({ + "file_name": "wan_ti2v_high_quality", + "base_model": "Wan 2.2", + "civitai": {}, + }) + assert set(result) == {"HIGH", "TI2V"} + + def test_file_name_lightning_turbo(self): + result = extract_auto_tags({ + "file_name": "sdxl_lightning_turbo_v3", + "base_model": "SDXL", + "civitai": {}, + }) + assert set(result) == {"Lightning", "Turbo"} + + def test_base_model_source(self): + result = extract_auto_tags({ + "file_name": "my_lora_v1", + "base_model": "Wan Video 2.2 I2V-A14B", + "civitai": {}, + }) + assert "I2V" in result + + def test_civitai_name_source(self): + result = extract_auto_tags({ + "file_name": "model_v1", + "base_model": "Wan", + "civitai": {"name": "HIGH Quality"}, + }) + assert "HIGH" in result + + def test_no_false_match_flow(self): + result = extract_auto_tags({ + "file_name": "flux_dev_model", + "base_model": "Flux.1 D", + "civitai": {}, + }) + assert "LOW" not in result + + def test_no_false_match_glow(self): + result = extract_auto_tags({ + "file_name": "glow_style_lora", + "base_model": "SDXL", + "civitai": {}, + }) + assert "LOW" not in result + + def test_high_low_only_for_wan(self): + """HIGH/LOW should not appear for non-Wan models even in filename.""" + result = extract_auto_tags({ + "file_name": "my_model_high_quality_v2", + "base_model": "Flux.1 D", + "civitai": {"name": "HIGH"}, + }) + assert "HIGH" not in result + assert "LOW" not in result + + def test_no_distilled(self): + result = extract_auto_tags({ + "file_name": "ltx-2.3-22b-distilled-lora-384", + "base_model": "LTXV 2.3", + "civitai": {}, + }) + assert result == [] + + def test_empty(self): + result = extract_auto_tags({ + "file_name": "generic_lora_v1", + "base_model": "SDXL", + "civitai": {}, + }) + assert result == [] + + def test_missing_fields(self): + result = extract_auto_tags({}) + assert result == [] + + def test_dash_separated(self): + result = extract_auto_tags({ + "file_name": "wan-i2v-high-v2", + "base_model": "Wan 2.2", + "civitai": {}, + }) + assert set(result) == {"HIGH", "I2V"} + + def test_dot_separated(self): + result = extract_auto_tags({ + "file_name": "wan.i2v.high.v2", + "base_model": "Wan 2.2", + "civitai": {}, + }) + assert set(result) == {"HIGH", "I2V"} + + def test_case_insensitive(self): + result = extract_auto_tags({ + "file_name": "WAN_i2v_High", + "base_model": "Wan 2.2", + "civitai": {}, + }) + assert set(result) == {"HIGH", "I2V"} + + +class TestAutoTagCategories: + def test_all_patterns_compile(self): + import re + for label, pattern in AUTO_TAG_CATEGORIES.items(): + re.compile(pattern, re.IGNORECASE) + + def test_mode_group_tags(self): + from services.auto_tag_service import MODE_TAGS + assert "HIGH" in MODE_TAGS + assert "LOW" in MODE_TAGS + + def test_video_group_tags(self): + from services.auto_tag_service import VIDEO_MODE_TAGS + assert "I2V" in VIDEO_MODE_TAGS + assert "T2V" in VIDEO_MODE_TAGS + assert "TI2V" in VIDEO_MODE_TAGS + + def test_default_enabled_groups(self): + from services.auto_tag_service import DEFAULT_ENABLED_GROUPS + assert "mode" in DEFAULT_ENABLED_GROUPS + assert "video" in DEFAULT_ENABLED_GROUPS + assert "speed" not in DEFAULT_ENABLED_GROUPS