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

refactor: Simplify model version existence checks and enhance version retrieval methods in scanners

Browse Source
This commit is contained in:
Will Miao
2025-07-09 10:26:03 +08:00
parent df9b554ce1
commit c692713ffb
3 changed files with 75 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
py/routes/misc_routes.py
View File

@@ -636,32 +636,18 @@ class MiscRoutes:
'error': 'Parameter modelVersionId must be an integer' 'error': 'Parameter modelVersionId must be an integer'
}, status=400) }, status=400)
# Check if the specific version exists in either scanner's cache # Check if the specific version exists in either scanner
exists = False exists = False
model_type = None model_type = None
# First check lora cache # Check lora scanner first
lora_cache = await lora_scanner.get_cached_data() if await lora_scanner.check_model_version_exists(model_id, model_version_id):
if lora_cache and lora_cache.raw_data: exists = True
for item in lora_cache.raw_data: model_type = 'lora'
if (item.get('civitai') and # If not found in lora, check checkpoint scanner
item['civitai'].get('modelId') == model_id and elif checkpoint_scanner and await checkpoint_scanner.check_model_version_exists(model_id, model_version_id):
item['civitai'].get('id') == model_version_id): exists = True
exists = True model_type = 'checkpoint'
model_type = 'lora'
break
# If not found in lora cache, check checkpoint cache
if not exists and checkpoint_scanner:
checkpoint_cache = await checkpoint_scanner.get_cached_data()
if checkpoint_cache and checkpoint_cache.raw_data:
for item in checkpoint_cache.raw_data:
if (item.get('civitai') and
item['civitai'].get('modelId') == model_id and
item['civitai'].get('id') == model_version_id):
exists = True
model_type = 'checkpoint'
break
return web.json_response({ return web.json_response({
'success': True, 'success': True,
@@ -671,35 +657,13 @@ class MiscRoutes:
# If modelVersionId is not provided, return all version IDs for the model # If modelVersionId is not provided, return all version IDs for the model
else: else:
# Lists to collect version IDs from both scanners # Get versions from lora scanner first
lora_versions = [] lora_versions = await lora_scanner.get_model_versions_by_id(model_id)
checkpoint_versions = [] checkpoint_versions = []
# Check lora cache # Only check checkpoint scanner if no lora versions found
lora_cache = await lora_scanner.get_cached_data() if not lora_versions:
if lora_cache and lora_cache.raw_data: checkpoint_versions = await checkpoint_scanner.get_model_versions_by_id(model_id)
for item in lora_cache.raw_data:
if (item.get('civitai') and
item['civitai'].get('modelId') == model_id and
item['civitai'].get('id')):
lora_versions.append({
'versionId': item['civitai'].get('id'),
'name': item['civitai'].get('name'),
'fileName': item.get('file_name', '')
})
# Check checkpoint cache
checkpoint_cache = await checkpoint_scanner.get_cached_data()
if checkpoint_cache and checkpoint_cache.raw_data:
for item in checkpoint_cache.raw_data:
if (item.get('civitai') and
item['civitai'].get('modelId') == model_id and
item['civitai'].get('id')):
checkpoint_versions.append({
'versionId': item['civitai'].get('id'),
'name': item['civitai'].get('name'),
'fileName': item.get('file_name', '')
})
# Determine model type and combine results # Determine model type and combine results
model_type = None model_type = None

5
py/services/download_manager.py
View File

@@ -82,6 +82,11 @@ class DownloadManager:
else: else:
return {'success': False, 'error': f'Model type "{model_type_from_info}" is not supported for download'} return {'success': False, 'error': f'Model type "{model_type_from_info}" is not supported for download'}
scanner = model_type == 'checkpoint' and await self._get_checkpoint_scanner() or await self._get_lora_scanner()
if scanner.check_model_version_exists(model_id, model_version_id):
return {'success': False, 'error': 'Model version already exists in library'}
# Handle use_default_paths # Handle use_default_paths
if use_default_paths: if use_default_paths:
# Set save_dir based on model type # Set save_dir based on model type

56
py/services/model_scanner.py
View File

@@ -1362,3 +1362,59 @@ class ModelScanner:
if file_name in self._hash_index._duplicate_filenames: if file_name in self._hash_index._duplicate_filenames:
if len(self._hash_index._duplicate_filenames[file_name]) <= 1: if len(self._hash_index._duplicate_filenames[file_name]) <= 1:
del self._hash_index._duplicate_filenames[file_name] del self._hash_index._duplicate_filenames[file_name]
async def check_model_version_exists(self, model_id: int, model_version_id: int) -> bool:
"""Check if a specific model version exists in the cache
Args:
model_id: Civitai model ID
model_version_id: Civitai model version ID
Returns:
bool: True if the model version exists, False otherwise
"""
try:
cache = await self.get_cached_data()
if not cache or not cache.raw_data:
return False
for item in cache.raw_data:
if (item.get('civitai') and
item['civitai'].get('modelId') == model_id and
item['civitai'].get('id') == model_version_id):
return True
return False
except Exception as e:
logger.error(f"Error checking model version existence: {e}")
return False
async def get_model_versions_by_id(self, model_id: int) -> List[Dict]:
"""Get all versions of a model by its ID
Args:
model_id: Civitai model ID
Returns:
List[Dict]: List of version information dictionaries
"""
try:
cache = await self.get_cached_data()
if not cache or not cache.raw_data:
return []
versions = []
for item in cache.raw_data:
if (item.get('civitai') and
item['civitai'].get('modelId') == model_id and
item['civitai'].get('id')):
versions.append({
'versionId': item['civitai'].get('id'),
'name': item['civitai'].get('name'),
'fileName': item.get('file_name', '')
})
return versions
except Exception as e:
logger.error(f"Error getting model versions: {e}")
return []