feat: add embedding management functionality with routes, services, and UI integration

This commit is contained in:
Will Miao
2025-07-25 21:14:56 +08:00
parent ea8a64fafc
commit c5a3af2399
10 changed files with 376 additions and 8 deletions

View File

@@ -0,0 +1,26 @@
import logging
from typing import List
from ..utils.models import EmbeddingMetadata
from ..config import config
from .model_scanner import ModelScanner
from .model_hash_index import ModelHashIndex
logger = logging.getLogger(__name__)
class EmbeddingScanner(ModelScanner):
"""Service for scanning and managing embedding files"""
def __init__(self):
# Define supported file extensions
file_extensions = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft'}
super().__init__(
model_type="embedding",
model_class=EmbeddingMetadata,
file_extensions=file_extensions,
hash_index=ModelHashIndex()
)
def get_model_roots(self) -> List[str]:
"""Get embedding root directories"""
return config.embeddings_roots

View File

@@ -0,0 +1,51 @@
import os
import logging
from typing import Dict, List, Optional
from .base_model_service import BaseModelService
from ..utils.models import EmbeddingMetadata
from ..config import config
from ..utils.routes_common import ModelRouteUtils
logger = logging.getLogger(__name__)
class EmbeddingService(BaseModelService):
"""Embedding-specific service implementation"""
def __init__(self, scanner):
"""Initialize Embedding service
Args:
scanner: Embedding scanner instance
"""
super().__init__("embedding", scanner, EmbeddingMetadata)
async def format_response(self, embedding_data: Dict) -> Dict:
"""Format Embedding data for API response"""
return {
"model_name": embedding_data["model_name"],
"file_name": embedding_data["file_name"],
"preview_url": config.get_preview_static_url(embedding_data.get("preview_url", "")),
"preview_nsfw_level": embedding_data.get("preview_nsfw_level", 0),
"base_model": embedding_data.get("base_model", ""),
"folder": embedding_data["folder"],
"sha256": embedding_data.get("sha256", ""),
"file_path": embedding_data["file_path"].replace(os.sep, "/"),
"file_size": embedding_data.get("size", 0),
"modified": embedding_data.get("modified", ""),
"tags": embedding_data.get("tags", []),
"modelDescription": embedding_data.get("modelDescription", ""),
"from_civitai": embedding_data.get("from_civitai", True),
"notes": embedding_data.get("notes", ""),
"model_type": embedding_data.get("model_type", "embedding"),
"favorite": embedding_data.get("favorite", False),
"civitai": ModelRouteUtils.filter_civitai_data(embedding_data.get("civitai", {}))
}
def find_duplicate_hashes(self) -> Dict:
"""Find Embeddings with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict:
"""Find Embeddings with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()

View File

@@ -122,11 +122,13 @@ class ModelServiceFactory:
def register_default_model_types():
"""Register the default model types (LoRA and Checkpoint)"""
"""Register the default model types (LoRA, Checkpoint, and Embedding)"""
from ..services.lora_service import LoraService
from ..services.checkpoint_service import CheckpointService
from ..services.embedding_service import EmbeddingService
from ..routes.lora_routes import LoraRoutes
from ..routes.checkpoint_routes import CheckpointRoutes
from ..routes.embedding_routes import EmbeddingRoutes
# Register LoRA model type
ModelServiceFactory.register_model_type('lora', LoraService, LoraRoutes)
@@ -134,4 +136,7 @@ def register_default_model_types():
# Register Checkpoint model type
ModelServiceFactory.register_model_type('checkpoint', CheckpointService, CheckpointRoutes)
logger.info("Registered default model types: lora, checkpoint")
# Register Embedding model type
ModelServiceFactory.register_model_type('embedding', EmbeddingService, EmbeddingRoutes)
logger.info("Registered default model types: lora, checkpoint, embedding")

View File

@@ -174,6 +174,27 @@ class ServiceRegistry:
logger.debug(f"Registered {service_name}")
return ws_manager
@classmethod
async def get_embedding_scanner(cls):
"""Get or create Embedding scanner instance"""
service_name = "embedding_scanner"
if service_name in cls._services:
return cls._services[service_name]
async with cls._get_lock(service_name):
# Double-check after acquiring lock
if service_name in cls._services:
return cls._services[service_name]
# Import here to avoid circular imports
from .embedding_scanner import EmbeddingScanner
scanner = await EmbeddingScanner.get_instance()
cls._services[service_name] = scanner
logger.debug(f"Created and registered {service_name}")
return scanner
@classmethod
def clear_services(cls):
"""Clear all registered services - mainly for testing"""