alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-24 14:42:11 -03:00
Code Issues Packages Projects Releases Wiki Activity

Refactor metadata handling to use unified provider system

Browse Source 
- Replaced direct usage of Civitai client with a fallback metadata provider across all recipe parsers.
- Updated metadata service to improve initialization and error handling.
- Enhanced download manager to utilize a downloader service for file operations.
- Improved recipe scanner to fetch model information through the new metadata provider.
- Updated utility functions to streamline image downloading and processing.
- Added comprehensive logging and error handling for better debugging and reliability.
- Introduced `get_default_metadata_provider()` for simplified access to the default provider.
- Ensured backward compatibility with existing APIs and workflows.
This commit is contained in:
Will Miao
2025-09-09 20:57:45 +08:00
parent 1ea468cfc4
commit 6fd74952b7
15 changed files with 350 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
py/services/download_manager.py
View File

@@ -10,6 +10,8 @@ from ..utils.exif_utils import ExifUtils
from ..utils.metadata_manager import MetadataManager
from .service_registry import ServiceRegistry
from .settings_manager import settings
from .metadata_service import get_default_metadata_provider
from .downloader import get_downloader
# Download to temporary file first
import tempfile
@@ -199,11 +201,11 @@ class DownloadManager:
if await embedding_scanner.check_model_version_exists(model_version_id):
return {'success': False, 'error': 'Model version already exists in embedding library'}
# Get civitai client
civitai_client = await self._get_civitai_client()
# Get metadata provider instead of civitai client directly
metadata_provider = await get_default_metadata_provider()
# Get version info based on the provided identifier
version_info = await civitai_client.get_model_version(model_id, model_version_id)
version_info = await metadata_provider.get_model_version(model_id, model_version_id)
if not version_info:
return {'success': False, 'error': 'Failed to fetch model metadata'}
@@ -445,8 +447,14 @@ class DownloadManager:
preview_ext = '.mp4'
preview_path = os.path.splitext(save_path)[0] + preview_ext
# Download video directly
if await civitai_client.download_preview_image(images[0]['url'], preview_path):
# Download video directly using downloader
downloader = await get_downloader()
success, result = await downloader.download_file(
images[0]['url'],
preview_path,
use_auth=False # Preview images typically don't need auth
)
if success:
metadata.preview_url = preview_path.replace(os.sep, '/')
metadata.preview_nsfw_level = images[0].get('nsfwLevel', 0)
else:
@@ -454,8 +462,16 @@ class DownloadManager:
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
temp_path = temp_file.name
# Download the original image to temp path
if await civitai_client.download_preview_image(images[0]['url'], temp_path):
# Download the original image to temp path using downloader
downloader = await get_downloader()
success, content = await downloader.download_to_memory(
images[0]['url'],
use_auth=False
)
if success:
# Save to temp file
with open(temp_path, 'wb') as f:
f.write(content)
# Optimize and convert to WebP
preview_path = os.path.splitext(save_path)[0] + '.webp'
@@ -486,12 +502,13 @@ class DownloadManager:
if progress_callback:
await progress_callback(3) # 3% progress after preview download
# Download model file with progress tracking
success, result = await civitai_client.download_file(
# Download model file with progress tracking using downloader
downloader = await get_downloader()
success, result = await downloader.download_file(
download_url,
save_dir,
os.path.basename(save_path),
progress_callback=lambda p: self._handle_download_progress(p, progress_callback)
save_path, # Use full path instead of separate dir and filename
progress_callback=lambda p: self._handle_download_progress(p, progress_callback),
use_auth=True # Model downloads need authentication
)
if not success:

4
py/services/downloader.py
View File

@@ -276,6 +276,10 @@ class Downloader:
while rename_attempt < max_rename_attempts and not rename_success:
try:
# If the destination file exists, remove it first (Windows safe)
if os.path.exists(save_path):
os.remove(save_path)
os.rename(part_path, save_path)
rename_success = True
except PermissionError as e:

68
py/services/metadata_service.py
View File

@@ -16,6 +16,10 @@ async def initialize_metadata_providers():
"""Initialize and configure all metadata providers based on settings"""
provider_manager = await ModelMetadataProviderManager.get_instance()
# Clear existing providers to allow reinitialization
provider_manager.providers.clear()
provider_manager.default_provider = None
# Get settings
enable_archive_db = settings.get('enable_metadata_archive_db', False)
priority = settings.get('metadata_provider_priority', 'archive_db')
@@ -24,23 +28,23 @@ async def initialize_metadata_providers():
# Initialize archive database provider if enabled
if enable_archive_db:
# Initialize archive manager
base_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
archive_manager = MetadataArchiveManager(base_path)
db_path = archive_manager.get_database_path()
if db_path:
try:
try:
# Initialize archive manager
base_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
archive_manager = MetadataArchiveManager(base_path)
db_path = archive_manager.get_database_path()
if db_path and os.path.exists(db_path):
sqlite_provider = SQLiteModelMetadataProvider(db_path)
provider_manager.register_provider('sqlite', sqlite_provider)
providers.append(('sqlite', sqlite_provider))
logger.info(f"SQLite metadata provider registered with database: {db_path}")
except Exception as e:
logger.error(f"Failed to initialize SQLite metadata provider: {e}")
else:
logger.warning("Metadata archive database is enabled but not available")
else:
logger.warning("Metadata archive database is enabled but database file not found")
except Exception as e:
logger.error(f"Failed to initialize SQLite metadata provider: {e}")
# Initialize Civitai API provider
# Initialize Civitai API provider (always available as fallback)
try:
civitai_client = await ServiceRegistry.get_civitai_client()
civitai_provider = CivitaiModelMetadataProvider(civitai_client)
@@ -50,42 +54,48 @@ async def initialize_metadata_providers():
except Exception as e:
logger.error(f"Failed to initialize Civitai API metadata provider: {e}")
# Set up fallback provider based on priority
# Set up fallback provider based on priority and available providers
if len(providers) > 1:
# Order providers based on priority setting
ordered_providers = []
if priority == 'archive_db':
# Archive DB first, then Civitai API
ordered_providers = [p[1] for p in providers if p[0] == 'sqlite'] + [p[1] for p in providers if p[0] == 'civitai_api']
ordered_providers = [p[1] for p in providers if p[0] == 'sqlite']
ordered_providers.extend([p[1] for p in providers if p[0] == 'civitai_api'])
else:
# Civitai API first, then Archive DB
ordered_providers = [p[1] for p in providers if p[0] == 'civitai_api'] + [p[1] for p in providers if p[0] == 'sqlite']
ordered_providers = [p[1] for p in providers if p[0] == 'civitai_api']
ordered_providers.extend([p[1] for p in providers if p[0] == 'sqlite'])
if ordered_providers:
fallback_provider = FallbackMetadataProvider(ordered_providers)
provider_manager.register_provider('fallback', fallback_provider, is_default=True)
logger.info(f"Fallback metadata provider registered with priority: {priority}")
logger.info(f"Fallback metadata provider registered with {len(ordered_providers)} providers, priority: {priority}")
elif len(providers) == 1:
# Only one provider available, set it as default
provider_name, provider = providers[0]
provider_manager.register_provider(provider_name, provider, is_default=True)
logger.info(f"Single metadata provider registered as default: {provider_name}")
else:
logger.warning("No metadata providers available")
logger.warning("No metadata providers available - this may cause metadata lookup failures")
return provider_manager
async def update_metadata_provider_priority():
"""Update metadata provider priority based on current settings"""
provider_manager = await ModelMetadataProviderManager.get_instance()
# Get current settings
enable_archive_db = settings.get('enable_metadata_archive_db', False)
priority = settings.get('metadata_provider_priority', 'archive_db')
# Rebuild providers with new priority
await initialize_metadata_providers()
logger.info(f"Updated metadata provider priority to: {priority}")
try:
# Get current settings
enable_archive_db = settings.get('enable_metadata_archive_db', False)
priority = settings.get('metadata_provider_priority', 'archive_db')
# Reinitialize all providers with new settings
provider_manager = await initialize_metadata_providers()
logger.info(f"Updated metadata provider priority to: {priority}, archive_db enabled: {enable_archive_db}")
return provider_manager
except Exception as e:
logger.error(f"Failed to update metadata provider priority: {e}")
return await ModelMetadataProviderManager.get_instance()
async def get_metadata_archive_manager():
"""Get metadata archive manager instance"""
@@ -100,3 +110,7 @@ async def get_metadata_provider(provider_name: str = None):
return provider_manager._get_provider(provider_name)
return provider_manager._get_provider()
async def get_default_metadata_provider():
"""Get the default metadata provider (fallback or single provider)"""
return await get_metadata_provider()

7
py/services/model_scanner.py
View File

@@ -730,11 +730,10 @@ class ModelScanner:
if needs_metadata_update and model_id:
logger.debug(f"Fetching missing metadata for {file_path} with model ID {model_id}")
from ..services.civitai_client import CivitaiClient
client = CivitaiClient()
from ..services.metadata_service import get_default_metadata_provider
metadata_provider = await get_default_metadata_provider()
model_metadata, status_code = await client.get_model_metadata(model_id)
await client.close()
model_metadata, status_code = await metadata_provider.get_model_metadata(model_id)
if status_code == 404:
logger.warning(f"Model {model_id} appears to be deleted from Civitai (404 response)")

11
py/services/recipe_scanner.py
View File

@@ -8,6 +8,7 @@ from ..config import config
from .recipe_cache import RecipeCache
from .service_registry import ServiceRegistry
from .lora_scanner import LoraScanner
from .metadata_service import get_default_metadata_provider
from ..utils.utils import fuzzy_match
from natsort import natsorted
import sys
@@ -431,13 +432,13 @@ class RecipeScanner:
async def _get_hash_from_civitai(self, model_version_id: str) -> Optional[str]:
"""Get hash from Civitai API"""
try:
# Get CivitaiClient from ServiceRegistry
civitai_client = await self._get_civitai_client()
if not civitai_client:
logger.error("Failed to get CivitaiClient from ServiceRegistry")
# Get metadata provider instead of civitai client directly
metadata_provider = await get_default_metadata_provider()
if not metadata_provider:
logger.error("Failed to get metadata provider")
return None
version_info, error_msg = await civitai_client.get_model_version_info(model_version_id)
version_info, error_msg = await metadata_provider.get_model_version_info(model_version_id)
if not version_info:
if error_msg and "model not found" in error_msg.lower():

2
py/services/settings_manager.py
View File

@@ -81,7 +81,7 @@ class SettingsManager:
return {
"civitai_api_key": "",
"show_only_sfw": False,
"language": "en", # 添加默认语言设置
"language": "en",
"enable_metadata_archive_db": False, # Enable metadata archive database
"metadata_provider_priority": "archive_db" # Default priority: 'archive_db' or 'civitai_api'
}