alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-21 21:22: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

119
METADATA_PROVIDER_REFACTOR_SUMMARY.md Normal file
View File

@@ -0,0 +1,119 @@
# Metadata Provider Refactor Summary
## Overview
This refactor improves the metadata provider initialization logic and replaces direct Civitai client usage with the unified FallbackMetadataProvider system throughout the codebase.
## Key Changes
### 1. Enhanced Metadata Service (`py/services/metadata_service.py`)
#### Improved `initialize_metadata_providers()`:
- Added provider clearing for proper reinitialization
- Enhanced error handling and validation
- Better logging for debugging
- Improved provider ordering logic based on priority settings
- More robust database path validation
#### Enhanced `update_metadata_provider_priority()`:
- More robust error handling
- Proper reinitalization of all providers
- Better logging for setting changes
#### New helper function:
- Added `get_default_metadata_provider()` for easier access to the default provider
### 2. Updated Recipe Parsers
All recipe parsers now use the unified metadata provider instead of direct civitai_client:
#### Files Updated:
- `py/recipes/parsers/civitai_image.py`
- `py/recipes/parsers/comfy.py`
- `py/recipes/parsers/automatic.py`
- `py/recipes/parsers/recipe_format.py`
- `py/recipes/parsers/meta_format.py`
#### Changes Made:
- Added import for `get_default_metadata_provider`
- Replaced `civitai_client.get_model_by_hash()` with `metadata_provider.get_model_by_hash()`
- Replaced `civitai_client.get_model_version_info()` with `metadata_provider.get_model_version_info()`
- Updated method signatures to indicate civitai_client parameter is deprecated
### 3. Download Manager Updates (`py/services/download_manager.py`)
#### Metadata Operations:
- Replaced direct civitai_client usage with metadata_provider for:
- `get_model_version()` calls for version info
#### Download Operations:
- Replaced `civitai_client.download_file()` with direct `downloader.download_file()` calls
- Replaced `civitai_client.download_preview_image()` with `downloader.download_to_memory()` for images
- Added proper authentication flags (`use_auth=True` for model files, `use_auth=False` for preview images)
### 4. Recipe Scanner Updates (`py/services/recipe_scanner.py`)
- Added import for `get_default_metadata_provider`
- Replaced `civitai_client.get_model_version_info()` with `metadata_provider.get_model_version_info()`
### 5. Utility Functions Updates (`py/utils/routes_common.py`)
- Added import for `get_downloader`
- Replaced preview image downloads with direct downloader usage
- Improved image optimization logic to work with in-memory downloads
- Better error handling for download and image processing operations
## Benefits
### 1. Unified Metadata Access
- All metadata requests now go through the fallback provider system
- Automatic failover between SQLite archive database and Civitai API
- Consistent metadata access patterns across all components
### 2. Improved Download Performance
- Direct use of the optimized downloader service
- Better connection pooling and retry logic
- Proper authentication handling
- Support for resumable downloads
### 3. Better Configuration Management
- Settings changes now properly update provider priority
- Clear separation between metadata and download operations
- Improved error handling and logging
### 4. Enhanced Reliability
- Fallback mechanisms ensure metadata is always available when possible
- Better error handling and recovery
- Consistent behavior across all parsers and services
## Usage
### Settings Changes
When users change metadata provider settings:
1. The `update_metadata_provider_priority()` function is automatically called
2. All providers are reinitialized with the new settings
3. The fallback provider is updated with the correct priority order
### Metadata Access
All components now use:
```python
from ...services.metadata_service import get_default_metadata_provider
metadata_provider = await get_default_metadata_provider()
result = await metadata_provider.get_model_by_hash(hash_value)
```
### Downloads
All downloads now use the unified downloader:
```python
from ...services.downloader import get_downloader
downloader = await get_downloader()
success, result = await downloader.download_file(url, path, use_auth=True)
```
## Compatibility
- All existing APIs and interfaces remain unchanged
- Backward compatibility maintained for existing workflows
- No changes required for external integrations
## Testing
- All updated files pass syntax validation
- Existing functionality preserved
- Enhanced error handling and logging for better debugging

12
py/recipes/parsers/automatic.py
View File

@@ -6,6 +6,7 @@ import logging
from typing import Dict, Any from typing import Dict, Any
from ..base import RecipeMetadataParser from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS from ..constants import GEN_PARAM_KEYS
from ...services.metadata_service import get_default_metadata_provider
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -30,6 +31,9 @@ class AutomaticMetadataParser(RecipeMetadataParser):
async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]: async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]:
"""Parse metadata from Automatic1111 format""" """Parse metadata from Automatic1111 format"""
try: try:
# Get metadata provider instead of using civitai_client directly
metadata_provider = await get_default_metadata_provider()
# Split on Negative prompt if it exists # Split on Negative prompt if it exists
if "Negative prompt:" in user_comment: if "Negative prompt:" in user_comment:
parts = user_comment.split('Negative prompt:', 1) parts = user_comment.split('Negative prompt:', 1)
@@ -216,9 +220,9 @@ class AutomaticMetadataParser(RecipeMetadataParser):
} }
# Get additional info from Civitai # Get additional info from Civitai
if civitai_client: if metadata_provider:
try: try:
civitai_info = await civitai_client.get_model_version_info(resource.get("modelVersionId")) civitai_info = await metadata_provider.get_model_version_info(resource.get("modelVersionId"))
populated_entry = await self.populate_lora_from_civitai( populated_entry = await self.populate_lora_from_civitai(
lora_entry, lora_entry,
civitai_info, civitai_info,
@@ -271,11 +275,11 @@ class AutomaticMetadataParser(RecipeMetadataParser):
} }
# Try to get info from Civitai # Try to get info from Civitai
if civitai_client: if metadata_provider:
try: try:
if lora_hash: if lora_hash:
# If we have hash, use it for lookup # If we have hash, use it for lookup
civitai_info = await civitai_client.get_model_by_hash(lora_hash) civitai_info = await metadata_provider.get_model_by_hash(lora_hash)
else: else:
civitai_info = None civitai_info = None

32
py/recipes/parsers/civitai_image.py
View File

@@ -5,6 +5,7 @@ import logging
from typing import Dict, Any, Union from typing import Dict, Any, Union
from ..base import RecipeMetadataParser from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS from ..constants import GEN_PARAM_KEYS
from ...services.metadata_service import get_default_metadata_provider
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -36,12 +37,15 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
Args: Args:
metadata: The metadata from the image (dict) metadata: The metadata from the image (dict)
recipe_scanner: Optional recipe scanner service recipe_scanner: Optional recipe scanner service
civitai_client: Optional Civitai API client civitai_client: Optional Civitai API client (deprecated, use metadata_provider instead)
Returns: Returns:
Dict containing parsed recipe data Dict containing parsed recipe data
""" """
try: try:
# Get metadata provider instead of using civitai_client directly
metadata_provider = await get_default_metadata_provider()
# Initialize result structure # Initialize result structure
result = { result = {
'base_model': None, 'base_model': None,
@@ -85,9 +89,9 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
# Extract base model information - directly if available # Extract base model information - directly if available
if "baseModel" in metadata: if "baseModel" in metadata:
result["base_model"] = metadata["baseModel"] result["base_model"] = metadata["baseModel"]
elif "Model hash" in metadata and civitai_client: elif "Model hash" in metadata and metadata_provider:
model_hash = metadata["Model hash"] model_hash = metadata["Model hash"]
model_info = await civitai_client.get_model_by_hash(model_hash) model_info = await metadata_provider.get_model_by_hash(model_hash)
if model_info: if model_info:
result["base_model"] = model_info.get("baseModel", "") result["base_model"] = model_info.get("baseModel", "")
elif "Model" in metadata and isinstance(metadata.get("resources"), list): elif "Model" in metadata and isinstance(metadata.get("resources"), list):
@@ -95,8 +99,8 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
for resource in metadata.get("resources", []): for resource in metadata.get("resources", []):
if resource.get("type") == "model" and resource.get("name") == metadata.get("Model"): if resource.get("type") == "model" and resource.get("name") == metadata.get("Model"):
# This is likely the checkpoint model # This is likely the checkpoint model
if civitai_client and resource.get("hash"): if metadata_provider and resource.get("hash"):
model_info = await civitai_client.get_model_by_hash(resource.get("hash")) model_info = await metadata_provider.get_model_by_hash(resource.get("hash"))
if model_info: if model_info:
result["base_model"] = model_info.get("baseModel", "") result["base_model"] = model_info.get("baseModel", "")
@@ -138,9 +142,9 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
} }
# Try to get info from Civitai if hash is available # Try to get info from Civitai if hash is available
if lora_entry['hash'] and civitai_client: if lora_entry['hash'] and metadata_provider:
try: try:
civitai_info = await civitai_client.get_model_by_hash(lora_hash) civitai_info = await metadata_provider.get_model_by_hash(lora_hash)
populated_entry = await self.populate_lora_from_civitai( populated_entry = await self.populate_lora_from_civitai(
lora_entry, lora_entry,
@@ -194,10 +198,10 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
} }
# Try to get info from Civitai if modelVersionId is available # Try to get info from Civitai if modelVersionId is available
if version_id and civitai_client: if version_id and metadata_provider:
try: try:
# Use get_model_version_info instead of get_model_version # Use get_model_version_info instead of get_model_version
civitai_info, error = await civitai_client.get_model_version_info(version_id) civitai_info, error = await metadata_provider.get_model_version_info(version_id)
if error: if error:
logger.warning(f"Error getting model version info: {error}") logger.warning(f"Error getting model version info: {error}")
@@ -259,11 +263,11 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
'isDeleted': False 'isDeleted': False
} }
# If we have a version ID and civitai client, try to get more info # If we have a version ID and metadata provider, try to get more info
if version_id and civitai_client: if version_id and metadata_provider:
try: try:
# Use get_model_version_info with the version ID # Use get_model_version_info with the version ID
civitai_info, error = await civitai_client.get_model_version_info(version_id) civitai_info, error = await metadata_provider.get_model_version_info(version_id)
if error: if error:
logger.warning(f"Error getting model version info: {error}") logger.warning(f"Error getting model version info: {error}")
@@ -316,9 +320,9 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
} }
# Try to get info from Civitai if hash is available # Try to get info from Civitai if hash is available
if lora_entry['hash'] and civitai_client: if lora_entry['hash'] and metadata_provider:
try: try:
civitai_info = await civitai_client.get_model_by_hash(lora_hash) civitai_info = await metadata_provider.get_model_by_hash(lora_hash)
populated_entry = await self.populate_lora_from_civitai( populated_entry = await self.populate_lora_from_civitai(
lora_entry, lora_entry,

14
py/recipes/parsers/comfy.py
View File

@@ -6,6 +6,7 @@ import logging
from typing import Dict, Any from typing import Dict, Any
from ..base import RecipeMetadataParser from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS from ..constants import GEN_PARAM_KEYS
from ...services.metadata_service import get_default_metadata_provider
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -26,6 +27,9 @@ class ComfyMetadataParser(RecipeMetadataParser):
async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]: async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]:
"""Parse metadata from Civitai ComfyUI metadata format""" """Parse metadata from Civitai ComfyUI metadata format"""
try: try:
# Get metadata provider instead of using civitai_client directly
metadata_provider = await get_default_metadata_provider()
data = json.loads(user_comment) data = json.loads(user_comment)
loras = [] loras = []
@@ -73,10 +77,10 @@ class ComfyMetadataParser(RecipeMetadataParser):
'isDeleted': False 'isDeleted': False
} }
# Get additional info from Civitai if client is available # Get additional info from Civitai if metadata provider is available
if civitai_client: if metadata_provider:
try: try:
civitai_info_tuple = await civitai_client.get_model_version_info(model_version_id) civitai_info_tuple = await metadata_provider.get_model_version_info(model_version_id)
# Populate lora entry with Civitai info # Populate lora entry with Civitai info
populated_entry = await self.populate_lora_from_civitai( populated_entry = await self.populate_lora_from_civitai(
lora_entry, lora_entry,
@@ -116,9 +120,9 @@ class ComfyMetadataParser(RecipeMetadataParser):
} }
# Get additional checkpoint info from Civitai # Get additional checkpoint info from Civitai
if civitai_client: if metadata_provider:
try: try:
civitai_info_tuple = await civitai_client.get_model_version_info(checkpoint_version_id) civitai_info_tuple = await metadata_provider.get_model_version_info(checkpoint_version_id)
civitai_info, _ = civitai_info_tuple if isinstance(civitai_info_tuple, tuple) else (civitai_info_tuple, None) civitai_info, _ = civitai_info_tuple if isinstance(civitai_info_tuple, tuple) else (civitai_info_tuple, None)
# Populate checkpoint with Civitai info # Populate checkpoint with Civitai info
checkpoint = await self.populate_checkpoint_from_civitai(checkpoint, civitai_info) checkpoint = await self.populate_checkpoint_from_civitai(checkpoint, civitai_info)

10
py/recipes/parsers/meta_format.py
View File

@@ -5,6 +5,7 @@ import logging
from typing import Dict, Any from typing import Dict, Any
from ..base import RecipeMetadataParser from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS from ..constants import GEN_PARAM_KEYS
from ...services.metadata_service import get_default_metadata_provider
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -18,8 +19,11 @@ class MetaFormatParser(RecipeMetadataParser):
return re.search(self.METADATA_MARKER, user_comment, re.IGNORECASE | re.DOTALL) is not None return re.search(self.METADATA_MARKER, user_comment, re.IGNORECASE | re.DOTALL) is not None
async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]: async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]:
"""Parse metadata from images with meta format metadata""" """Parse metadata from images with meta format metadata (Lora_N Model hash format)"""
try: try:
# Get metadata provider instead of using civitai_client directly
metadata_provider = await get_default_metadata_provider()
# Extract prompt and negative prompt # Extract prompt and negative prompt
parts = user_comment.split('Negative prompt:', 1) parts = user_comment.split('Negative prompt:', 1)
prompt = parts[0].strip() prompt = parts[0].strip()
@@ -122,9 +126,9 @@ class MetaFormatParser(RecipeMetadataParser):
} }
# Get info from Civitai by hash if available # Get info from Civitai by hash if available
if civitai_client and hash_value: if metadata_provider and hash_value:
try: try:
civitai_info = await civitai_client.get_model_by_hash(hash_value) civitai_info = await metadata_provider.get_model_by_hash(hash_value)
# Populate lora entry with Civitai info # Populate lora entry with Civitai info
populated_entry = await self.populate_lora_from_civitai( populated_entry = await self.populate_lora_from_civitai(
lora_entry, lora_entry,

8
py/recipes/parsers/recipe_format.py
View File

@@ -7,6 +7,7 @@ from typing import Dict, Any
from ...config import config from ...config import config
from ..base import RecipeMetadataParser from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS from ..constants import GEN_PARAM_KEYS
from ...services.metadata_service import get_default_metadata_provider
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -23,6 +24,9 @@ class RecipeFormatParser(RecipeMetadataParser):
async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]: async def parse_metadata(self, user_comment: str, recipe_scanner=None, civitai_client=None) -> Dict[str, Any]:
"""Parse metadata from images with dedicated recipe metadata format""" """Parse metadata from images with dedicated recipe metadata format"""
try: try:
# Get metadata provider instead of using civitai_client directly
metadata_provider = await get_default_metadata_provider()
# Extract recipe metadata from user comment # Extract recipe metadata from user comment
try: try:
# Look for recipe metadata section # Look for recipe metadata section
@@ -71,9 +75,9 @@ class RecipeFormatParser(RecipeMetadataParser):
lora_entry['localPath'] = None lora_entry['localPath'] = None
# Try to get additional info from Civitai if we have a model version ID # Try to get additional info from Civitai if we have a model version ID
if lora.get('modelVersionId') and civitai_client: if lora.get('modelVersionId') and metadata_provider:
try: try:
civitai_info_tuple = await civitai_client.get_model_version_info(lora['modelVersionId']) civitai_info_tuple = await metadata_provider.get_model_version_info(lora['modelVersionId'])
# Populate lora entry with Civitai info # Populate lora entry with Civitai info
populated_entry = await self.populate_lora_from_civitai( populated_entry = await self.populate_lora_from_civitai(
lora_entry, lora_entry,

31
py/routes/recipe_routes.py
View File

@@ -24,6 +24,7 @@ from ..config import config
standalone_mode = 'nodes' not in sys.modules standalone_mode = 'nodes' not in sys.modules
from ..services.service_registry import ServiceRegistry # Add ServiceRegistry import from ..services.service_registry import ServiceRegistry # Add ServiceRegistry import
from ..services.downloader import get_downloader
# Only import MetadataRegistry in non-standalone mode # Only import MetadataRegistry in non-standalone mode
if not standalone_mode: if not standalone_mode:
@@ -372,21 +373,23 @@ class RecipeRoutes:
"loras": [] "loras": []
}, status=400) }, status=400)
# Download image directly from URL # Download image using unified downloader
session = await self.civitai_client.session downloader = await get_downloader()
# Create a temporary file to save the downloaded image # Create a temporary file to save the downloaded image
with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file: with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
temp_path = temp_file.name temp_path = temp_file.name
async with session.get(image_url) as response: success, result = await downloader.download_file(
if response.status != 200: image_url,
return web.json_response({ temp_path,
"error": f"Failed to download image from URL: HTTP {response.status}", use_auth=False # Image downloads typically don't need auth
"loras": [] )
}, status=400)
if not success:
with open(temp_path, 'wb') as f: return web.json_response({
f.write(await response.read()) "error": f"Failed to download image from URL: {result}",
"loras": []
}, status=400)
# Use meta field from image_info as metadata # Use meta field from image_info as metadata
if 'meta' in image_info: if 'meta' in image_info:
@@ -430,8 +433,7 @@ class RecipeRoutes:
# Parse the metadata # Parse the metadata
result = await parser.parse_metadata( result = await parser.parse_metadata(
metadata, metadata,
recipe_scanner=self.recipe_scanner, recipe_scanner=self.recipe_scanner
civitai_client=self.civitai_client
) )
# For URL mode, include the image data as base64 # For URL mode, include the image data as base64
@@ -532,8 +534,7 @@ class RecipeRoutes:
# Parse the metadata # Parse the metadata
result = await parser.parse_metadata( result = await parser.parse_metadata(
metadata, metadata,
recipe_scanner=self.recipe_scanner, recipe_scanner=self.recipe_scanner
civitai_client=self.civitai_client
) )
# Add base64 image data to result # Add base64 image data to result

4
py/routes/update_routes.py
View File

@@ -258,7 +258,7 @@ class UpdateRoutes:
try: try:
downloader = await Downloader.get_instance() downloader = await Downloader.get_instance()
success, data = await downloader.make_request('GET', github_url, headers={'Accept': 'application/vnd.github+json'}) success, data = await downloader.make_request('GET', github_url, custom_headers={'Accept': 'application/vnd.github+json'})
if not success: if not success:
logger.warning(f"Failed to fetch GitHub commit: {data}") logger.warning(f"Failed to fetch GitHub commit: {data}")
@@ -424,7 +424,7 @@ class UpdateRoutes:
try: try:
downloader = await Downloader.get_instance() downloader = await Downloader.get_instance()
success, data = await downloader.make_request('GET', github_url, headers={'Accept': 'application/vnd.github+json'}) success, data = await downloader.make_request('GET', github_url, custom_headers={'Accept': 'application/vnd.github+json'})
if not success: if not success:
logger.warning(f"Failed to fetch GitHub release: {data}") logger.warning(f"Failed to fetch GitHub release: {data}")

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 ..utils.metadata_manager import MetadataManager
from .service_registry import ServiceRegistry from .service_registry import ServiceRegistry
from .settings_manager import settings from .settings_manager import settings
from .metadata_service import get_default_metadata_provider
from .downloader import get_downloader
# Download to temporary file first # Download to temporary file first
import tempfile import tempfile
@@ -199,11 +201,11 @@ class DownloadManager:
if await embedding_scanner.check_model_version_exists(model_version_id): if await embedding_scanner.check_model_version_exists(model_version_id):
return {'success': False, 'error': 'Model version already exists in embedding library'} return {'success': False, 'error': 'Model version already exists in embedding library'}
# Get civitai client # Get metadata provider instead of civitai client directly
civitai_client = await self._get_civitai_client() metadata_provider = await get_default_metadata_provider()
# Get version info based on the provided identifier # 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: if not version_info:
return {'success': False, 'error': 'Failed to fetch model metadata'} return {'success': False, 'error': 'Failed to fetch model metadata'}
@@ -445,8 +447,14 @@ class DownloadManager:
preview_ext = '.mp4' preview_ext = '.mp4'
preview_path = os.path.splitext(save_path)[0] + preview_ext preview_path = os.path.splitext(save_path)[0] + preview_ext
# Download video directly # Download video directly using downloader
if await civitai_client.download_preview_image(images[0]['url'], preview_path): 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_url = preview_path.replace(os.sep, '/')
metadata.preview_nsfw_level = images[0].get('nsfwLevel', 0) metadata.preview_nsfw_level = images[0].get('nsfwLevel', 0)
else: else:
@@ -454,8 +462,16 @@ class DownloadManager:
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file: with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
temp_path = temp_file.name temp_path = temp_file.name
# Download the original image to temp path # Download the original image to temp path using downloader
if await civitai_client.download_preview_image(images[0]['url'], temp_path): 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 # Optimize and convert to WebP
preview_path = os.path.splitext(save_path)[0] + '.webp' preview_path = os.path.splitext(save_path)[0] + '.webp'
@@ -486,12 +502,13 @@ class DownloadManager:
if progress_callback: if progress_callback:
await progress_callback(3) # 3% progress after preview download await progress_callback(3) # 3% progress after preview download
# Download model file with progress tracking # Download model file with progress tracking using downloader
success, result = await civitai_client.download_file( downloader = await get_downloader()
success, result = await downloader.download_file(
download_url, download_url,
save_dir, save_path, # Use full path instead of separate dir and filename
os.path.basename(save_path), progress_callback=lambda p: self._handle_download_progress(p, progress_callback),
progress_callback=lambda p: self._handle_download_progress(p, progress_callback) use_auth=True # Model downloads need authentication
) )
if not success: 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: while rename_attempt < max_rename_attempts and not rename_success:
try: 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) os.rename(part_path, save_path)
rename_success = True rename_success = True
except PermissionError as e: 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""" """Initialize and configure all metadata providers based on settings"""
provider_manager = await ModelMetadataProviderManager.get_instance() provider_manager = await ModelMetadataProviderManager.get_instance()
# Clear existing providers to allow reinitialization
provider_manager.providers.clear()
provider_manager.default_provider = None
# Get settings # Get settings
enable_archive_db = settings.get('enable_metadata_archive_db', False) enable_archive_db = settings.get('enable_metadata_archive_db', False)
priority = settings.get('metadata_provider_priority', 'archive_db') priority = settings.get('metadata_provider_priority', 'archive_db')
@@ -24,23 +28,23 @@ async def initialize_metadata_providers():
# Initialize archive database provider if enabled # Initialize archive database provider if enabled
if enable_archive_db: if enable_archive_db:
# Initialize archive manager try:
base_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) # Initialize archive manager
archive_manager = MetadataArchiveManager(base_path) 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: db_path = archive_manager.get_database_path()
try: if db_path and os.path.exists(db_path):
sqlite_provider = SQLiteModelMetadataProvider(db_path) sqlite_provider = SQLiteModelMetadataProvider(db_path)
provider_manager.register_provider('sqlite', sqlite_provider) provider_manager.register_provider('sqlite', sqlite_provider)
providers.append(('sqlite', sqlite_provider)) providers.append(('sqlite', sqlite_provider))
logger.info(f"SQLite metadata provider registered with database: {db_path}") logger.info(f"SQLite metadata provider registered with database: {db_path}")
except Exception as e: else:
logger.error(f"Failed to initialize SQLite metadata provider: {e}") logger.warning("Metadata archive database is enabled but database file not found")
else: except Exception as e:
logger.warning("Metadata archive database is enabled but not available") logger.error(f"Failed to initialize SQLite metadata provider: {e}")
# Initialize Civitai API provider # Initialize Civitai API provider (always available as fallback)
try: try:
civitai_client = await ServiceRegistry.get_civitai_client() civitai_client = await ServiceRegistry.get_civitai_client()
civitai_provider = CivitaiModelMetadataProvider(civitai_client) civitai_provider = CivitaiModelMetadataProvider(civitai_client)
@@ -50,42 +54,48 @@ async def initialize_metadata_providers():
except Exception as e: except Exception as e:
logger.error(f"Failed to initialize Civitai API metadata provider: {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: if len(providers) > 1:
# Order providers based on priority setting # Order providers based on priority setting
ordered_providers = []
if priority == 'archive_db': if priority == 'archive_db':
# Archive DB first, then Civitai API # 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: else:
# Civitai API first, then Archive DB # 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: if ordered_providers:
fallback_provider = FallbackMetadataProvider(ordered_providers) fallback_provider = FallbackMetadataProvider(ordered_providers)
provider_manager.register_provider('fallback', fallback_provider, is_default=True) 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: elif len(providers) == 1:
# Only one provider available, set it as default # Only one provider available, set it as default
provider_name, provider = providers[0] provider_name, provider = providers[0]
provider_manager.register_provider(provider_name, provider, is_default=True) provider_manager.register_provider(provider_name, provider, is_default=True)
logger.info(f"Single metadata provider registered as default: {provider_name}") logger.info(f"Single metadata provider registered as default: {provider_name}")
else: else:
logger.warning("No metadata providers available") logger.warning("No metadata providers available - this may cause metadata lookup failures")
return provider_manager return provider_manager
async def update_metadata_provider_priority(): async def update_metadata_provider_priority():
"""Update metadata provider priority based on current settings""" """Update metadata provider priority based on current settings"""
provider_manager = await ModelMetadataProviderManager.get_instance() try:
# Get current settings
# Get current settings enable_archive_db = settings.get('enable_metadata_archive_db', False)
enable_archive_db = settings.get('enable_metadata_archive_db', False) priority = settings.get('metadata_provider_priority', 'archive_db')
priority = settings.get('metadata_provider_priority', 'archive_db')
# Reinitialize all providers with new settings
# Rebuild providers with new priority provider_manager = await initialize_metadata_providers()
await initialize_metadata_providers()
logger.info(f"Updated metadata provider priority to: {priority}, archive_db enabled: {enable_archive_db}")
logger.info(f"Updated metadata provider priority to: {priority}") 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(): async def get_metadata_archive_manager():
"""Get metadata archive manager instance""" """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(provider_name)
return provider_manager._get_provider() 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: if needs_metadata_update and model_id:
logger.debug(f"Fetching missing metadata for {file_path} with model ID {model_id}") logger.debug(f"Fetching missing metadata for {file_path} with model ID {model_id}")
from ..services.civitai_client import CivitaiClient from ..services.metadata_service import get_default_metadata_provider
client = CivitaiClient() metadata_provider = await get_default_metadata_provider()
model_metadata, status_code = await client.get_model_metadata(model_id) model_metadata, status_code = await metadata_provider.get_model_metadata(model_id)
await client.close()
if status_code == 404: if status_code == 404:
logger.warning(f"Model {model_id} appears to be deleted from Civitai (404 response)") 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 .recipe_cache import RecipeCache
from .service_registry import ServiceRegistry from .service_registry import ServiceRegistry
from .lora_scanner import LoraScanner from .lora_scanner import LoraScanner
from .metadata_service import get_default_metadata_provider
from ..utils.utils import fuzzy_match from ..utils.utils import fuzzy_match
from natsort import natsorted from natsort import natsorted
import sys import sys
@@ -431,13 +432,13 @@ class RecipeScanner:
async def _get_hash_from_civitai(self, model_version_id: str) -> Optional[str]: async def _get_hash_from_civitai(self, model_version_id: str) -> Optional[str]:
"""Get hash from Civitai API""" """Get hash from Civitai API"""
try: try:
# Get CivitaiClient from ServiceRegistry # Get metadata provider instead of civitai client directly
civitai_client = await self._get_civitai_client() metadata_provider = await get_default_metadata_provider()
if not civitai_client: if not metadata_provider:
logger.error("Failed to get CivitaiClient from ServiceRegistry") logger.error("Failed to get metadata provider")
return None 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 not version_info:
if error_msg and "model not found" in error_msg.lower(): 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 { return {
"civitai_api_key": "", "civitai_api_key": "",
"show_only_sfw": False, "show_only_sfw": False,
"language": "en", # 添加默认语言设置 "language": "en",
"enable_metadata_archive_db": False, # Enable metadata archive database "enable_metadata_archive_db": False, # Enable metadata archive database
"metadata_provider_priority": "archive_db" # Default priority: 'archive_db' or 'civitai_api' "metadata_provider_priority": "archive_db" # Default priority: 'archive_db' or 'civitai_api'
} }

165
py/utils/routes_common.py
View File

@@ -7,13 +7,12 @@ from aiohttp import web
from .model_utils import determine_base_model from .model_utils import determine_base_model
from .constants import PREVIEW_EXTENSIONS, CARD_PREVIEW_WIDTH from .constants import PREVIEW_EXTENSIONS, CARD_PREVIEW_WIDTH
from ..config import config from ..config import config
from ..services.civitai_client import CivitaiClient
from ..services.service_registry import ServiceRegistry from ..services.service_registry import ServiceRegistry
from ..services.downloader import get_downloader
from ..utils.exif_utils import ExifUtils from ..utils.exif_utils import ExifUtils
from ..utils.metadata_manager import MetadataManager from ..utils.metadata_manager import MetadataManager
from ..services.download_manager import DownloadManager
from ..services.websocket_manager import ws_manager from ..services.websocket_manager import ws_manager
from ..services.metadata_service import get_metadata_provider from ..services.metadata_service import get_default_metadata_provider
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -40,7 +39,7 @@ class ModelRouteUtils:
@staticmethod @staticmethod
async def update_model_metadata(metadata_path: str, local_metadata: Dict, async def update_model_metadata(metadata_path: str, local_metadata: Dict,
civitai_metadata: Dict, client: CivitaiClient) -> None: civitai_metadata: Dict, metadata_provider=None) -> None:
"""Update local metadata with CivitAI data""" """Update local metadata with CivitAI data"""
# Save existing trainedWords and customImages if they exist # Save existing trainedWords and customImages if they exist
existing_civitai = local_metadata.get('civitai') or {} # Use empty dict if None existing_civitai = local_metadata.get('civitai') or {} # Use empty dict if None
@@ -80,15 +79,17 @@ class ModelRouteUtils:
# If we have modelId and don't have enough metadata, fetch additional data # If we have modelId and don't have enough metadata, fetch additional data
if not model_metadata or not model_metadata.get('description'): if not model_metadata or not model_metadata.get('description'):
model_id = civitai_metadata.get('modelId') model_id = civitai_metadata.get('modelId')
if model_id: if model_id and metadata_provider:
fetched_metadata, _ = await client.get_model_metadata(str(model_id)) fetched_metadata, _ = await metadata_provider.get_model_metadata(str(model_id))
if fetched_metadata: if fetched_metadata:
model_metadata = fetched_metadata model_metadata = fetched_metadata
# Update local metadata with the model information # Update local metadata with the model information
if model_metadata: if model_metadata:
local_metadata['modelDescription'] = model_metadata.get('description', '') local_metadata['modelDescription'] = model_metadata.get('description', '')
local_metadata['tags'] = model_metadata.get('tags', []) # Only set tags if local_metadata['tags'] is empty
if not local_metadata.get('tags'):
local_metadata['tags'] = model_metadata.get('tags', [])
if 'creator' in model_metadata and model_metadata['creator']: if 'creator' in model_metadata and model_metadata['creator']:
local_metadata['civitai']['creator'] = model_metadata['creator'] local_metadata['civitai']['creator'] = model_metadata['creator']
@@ -114,22 +115,28 @@ class ModelRouteUtils:
preview_path = os.path.join(os.path.dirname(metadata_path), preview_filename) preview_path = os.path.join(os.path.dirname(metadata_path), preview_filename)
if is_video: if is_video:
# Download video as is # Download video as is using downloader
if await client.download_preview_image(first_preview['url'], preview_path): downloader = await get_downloader()
success, result = await downloader.download_file(
first_preview['url'],
preview_path,
use_auth=False
)
if success:
local_metadata['preview_url'] = preview_path.replace(os.sep, '/') local_metadata['preview_url'] = preview_path.replace(os.sep, '/')
local_metadata['preview_nsfw_level'] = first_preview.get('nsfwLevel', 0) local_metadata['preview_nsfw_level'] = first_preview.get('nsfwLevel', 0)
else: else:
# For images, download and then optimize to WebP # For images, download and then optimize to WebP using downloader
temp_path = preview_path + ".temp" downloader = await get_downloader()
if await client.download_preview_image(first_preview['url'], temp_path): success, content = await downloader.download_to_memory(
first_preview['url'],
use_auth=False
)
if success:
try: try:
# Read the downloaded image
with open(temp_path, 'rb') as f:
image_data = f.read()
# Optimize and convert to WebP # Optimize and convert to WebP
optimized_data, _ = ExifUtils.optimize_image( optimized_data, _ = ExifUtils.optimize_image(
image_data=image_data, image_data=content, # Use downloaded content directly
target_width=CARD_PREVIEW_WIDTH, target_width=CARD_PREVIEW_WIDTH,
format='webp', format='webp',
quality=85, quality=85,
@@ -144,17 +151,16 @@ class ModelRouteUtils:
local_metadata['preview_url'] = preview_path.replace(os.sep, '/') local_metadata['preview_url'] = preview_path.replace(os.sep, '/')
local_metadata['preview_nsfw_level'] = first_preview.get('nsfwLevel', 0) local_metadata['preview_nsfw_level'] = first_preview.get('nsfwLevel', 0)
# Remove the temporary file
if os.path.exists(temp_path):
os.remove(temp_path)
except Exception as e: except Exception as e:
logger.error(f"Error optimizing preview image: {e}") logger.error(f"Error optimizing preview image: {e}")
# If optimization fails, try to use the downloaded image directly # If optimization fails, save the original content
if os.path.exists(temp_path): try:
os.rename(temp_path, preview_path) with open(preview_path, 'wb') as f:
f.write(content)
local_metadata['preview_url'] = preview_path.replace(os.sep, '/') local_metadata['preview_url'] = preview_path.replace(os.sep, '/')
local_metadata['preview_nsfw_level'] = first_preview.get('nsfwLevel', 0) local_metadata['preview_nsfw_level'] = first_preview.get('nsfwLevel', 0)
except Exception as save_error:
logger.error(f"Error saving preview image: {save_error}")
# Save updated metadata # Save updated metadata
await MetadataManager.save_metadata(metadata_path, local_metadata) await MetadataManager.save_metadata(metadata_path, local_metadata)
@@ -177,7 +183,6 @@ class ModelRouteUtils:
Returns: Returns:
bool: True if successful, False otherwise bool: True if successful, False otherwise
""" """
client = CivitaiClient()
try: try:
# Validate input parameters # Validate input parameters
if not isinstance(model_data, dict): if not isinstance(model_data, dict):
@@ -189,8 +194,9 @@ class ModelRouteUtils:
# Check if model metadata exists # Check if model metadata exists
local_metadata = await ModelRouteUtils.load_local_metadata(metadata_path) local_metadata = await ModelRouteUtils.load_local_metadata(metadata_path)
# Fetch metadata from Civitai # Get metadata provider and fetch metadata from unified provider
civitai_metadata = await client.get_model_by_hash(sha256) metadata_provider = await get_default_metadata_provider()
civitai_metadata = await metadata_provider.get_model_by_hash(sha256)
if not civitai_metadata: if not civitai_metadata:
# Mark as not from CivitAI if not found # Mark as not from CivitAI if not found
local_metadata['from_civitai'] = False local_metadata['from_civitai'] = False
@@ -203,7 +209,7 @@ class ModelRouteUtils:
metadata_path, metadata_path,
local_metadata, local_metadata,
civitai_metadata, civitai_metadata,
client metadata_provider
) )
# Update cache object directly using safe .get() method # Update cache object directly using safe .get() method
@@ -226,8 +232,6 @@ class ModelRouteUtils:
except Exception as e: except Exception as e:
logger.error(f"Error fetching CivitAI data: {str(e)}", exc_info=True) # Include stack trace logger.error(f"Error fetching CivitAI data: {str(e)}", exc_info=True) # Include stack trace
return False return False
finally:
await client.close()
@staticmethod @staticmethod
def filter_civitai_data(data: Dict, minimal: bool = False) -> Dict: def filter_civitai_data(data: Dict, minimal: bool = False) -> Dict:
@@ -360,24 +364,22 @@ class ModelRouteUtils:
if not local_metadata or not local_metadata.get('sha256'): if not local_metadata or not local_metadata.get('sha256'):
return web.json_response({"success": False, "error": "No SHA256 hash found"}, status=400) return web.json_response({"success": False, "error": "No SHA256 hash found"}, status=400)
# Create a client for fetching from Civitai # Get metadata provider and fetch from unified provider
client = CivitaiClient() metadata_provider = await get_default_metadata_provider()
try:
# Fetch and update metadata # Fetch and update metadata
civitai_metadata = await client.get_model_by_hash(local_metadata["sha256"]) civitai_metadata = await metadata_provider.get_model_by_hash(local_metadata["sha256"])
if not civitai_metadata: if not civitai_metadata:
await ModelRouteUtils.handle_not_found_on_civitai(metadata_path, local_metadata) await ModelRouteUtils.handle_not_found_on_civitai(metadata_path, local_metadata)
return web.json_response({"success": False, "error": "Not found on CivitAI"}, status=404) return web.json_response({"success": False, "error": "Not found on CivitAI"}, status=404)
await ModelRouteUtils.update_model_metadata(metadata_path, local_metadata, civitai_metadata, client) await ModelRouteUtils.update_model_metadata(metadata_path, local_metadata, civitai_metadata, metadata_provider)
# Update the cache # Update the cache
await scanner.update_single_model_cache(data['file_path'], data['file_path'], local_metadata) await scanner.update_single_model_cache(data['file_path'], data['file_path'], local_metadata)
# Return the updated metadata along with success status # Return the updated metadata along with success status
return web.json_response({"success": True, "metadata": local_metadata}) return web.json_response({"success": True, "metadata": local_metadata})
finally:
await client.close()
except Exception as e: except Exception as e:
logger.error(f"Error fetching from CivitAI: {e}", exc_info=True) logger.error(f"Error fetching from CivitAI: {e}", exc_info=True)
@@ -778,43 +780,38 @@ class ModelRouteUtils:
# Check if model metadata exists # Check if model metadata exists
local_metadata = await ModelRouteUtils.load_local_metadata(metadata_path) local_metadata = await ModelRouteUtils.load_local_metadata(metadata_path)
# Create a client for fetching from Civitai # Get metadata provider and fetch metadata using get_model_version which includes more comprehensive data
client = await CivitaiClient.get_instance() metadata_provider = await get_default_metadata_provider()
try: civitai_metadata = await metadata_provider.get_model_version(model_id, model_version_id)
# Fetch metadata using get_model_version which includes more comprehensive data if not civitai_metadata:
civitai_metadata = await client.get_model_version(model_id, model_version_id) error_msg = f"Model version not found on CivitAI for ID: {model_id}"
if not civitai_metadata: if model_version_id:
error_msg = f"Model version not found on CivitAI for ID: {model_id}" error_msg += f" with version: {model_version_id}"
if model_version_id: return web.json_response({"success": False, "error": error_msg}, status=404)
error_msg += f" with version: {model_version_id}"
return web.json_response({"success": False, "error": error_msg}, status=404) # Try to find the primary model file to get the SHA256 hash
primary_model_file = None
# Try to find the primary model file to get the SHA256 hash for file in civitai_metadata.get('files', []):
primary_model_file = None if file.get('primary', False) and file.get('type') == 'Model':
for file in civitai_metadata.get('files', []): primary_model_file = file
if file.get('primary', False) and file.get('type') == 'Model': break
primary_model_file = file
break # Update the SHA256 hash in local metadata if available
if primary_model_file and primary_model_file.get('hashes', {}).get('SHA256'):
# Update the SHA256 hash in local metadata if available local_metadata['sha256'] = primary_model_file['hashes']['SHA256'].lower()
if primary_model_file and primary_model_file.get('hashes', {}).get('SHA256'):
local_metadata['sha256'] = primary_model_file['hashes']['SHA256'].lower() # Update metadata with CivitAI information
await ModelRouteUtils.update_model_metadata(metadata_path, local_metadata, civitai_metadata, metadata_provider)
# Update metadata with CivitAI information
await ModelRouteUtils.update_model_metadata(metadata_path, local_metadata, civitai_metadata, client) # Update the cache
await scanner.update_single_model_cache(file_path, file_path, local_metadata)
# Update the cache
await scanner.update_single_model_cache(file_path, file_path, local_metadata) return web.json_response({
"success": True,
return web.json_response({ "message": f"Model successfully re-linked to Civitai model {model_id}" +
"success": True, (f" version {model_version_id}" if model_version_id else ""),
"message": f"Model successfully re-linked to Civitai model {model_id}" + "hash": local_metadata.get('sha256', '')
(f" version {model_version_id}" if model_version_id else ""), })
"hash": local_metadata.get('sha256', '')
})
finally:
await client.close()
except Exception as e: except Exception as e:
logger.error(f"Error re-linking to CivitAI: {e}", exc_info=True) logger.error(f"Error re-linking to CivitAI: {e}", exc_info=True)