mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
84c62f2954
- Remove backward compatibility code for `model_type` in `ModelScanner._build_cache_entry()` - Update `CheckpointScanner` to only handle `sub_type` in `adjust_metadata()` and `adjust_cached_entry()` - Delete deprecated aliases `resolve_civitai_model_type` and `normalize_civitai_model_type` from `model_query.py` - Update frontend components (`RecipeModal.js`, `ModelCard.js`, etc.) to use `sub_type` instead of `model_type` - Update API response format to return only `sub_type`, removing `model_type` from service responses - Revise technical documentation to mark Phase 5 as completed and remove outdated TODO items All cleanup tasks for the model type refactoring are now complete, ensuring consistent use of `sub_type` across the codebase.
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
import logging
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from ..utils.models import CheckpointMetadata
|
|
from ..config import config
|
|
from .model_scanner import ModelScanner
|
|
from .model_hash_index import ModelHashIndex
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class CheckpointScanner(ModelScanner):
|
|
"""Service for scanning and managing checkpoint files"""
|
|
|
|
def __init__(self):
|
|
# Define supported file extensions
|
|
file_extensions = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft', '.gguf'}
|
|
super().__init__(
|
|
model_type="checkpoint",
|
|
model_class=CheckpointMetadata,
|
|
file_extensions=file_extensions,
|
|
hash_index=ModelHashIndex()
|
|
)
|
|
|
|
def _resolve_sub_type(self, root_path: Optional[str]) -> Optional[str]:
|
|
"""Resolve the sub-type based on the root path."""
|
|
if not root_path:
|
|
return None
|
|
|
|
if config.checkpoints_roots and root_path in config.checkpoints_roots:
|
|
return "checkpoint"
|
|
|
|
if config.unet_roots and root_path in config.unet_roots:
|
|
return "diffusion_model"
|
|
|
|
return None
|
|
|
|
def adjust_metadata(self, metadata, file_path, root_path):
|
|
"""Adjust metadata during scanning to set sub_type."""
|
|
sub_type = self._resolve_sub_type(root_path)
|
|
if sub_type:
|
|
metadata.sub_type = sub_type
|
|
return metadata
|
|
|
|
def adjust_cached_entry(self, entry: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Adjust entries loaded from the persisted cache to ensure sub_type is set."""
|
|
sub_type = self._resolve_sub_type(
|
|
self._find_root_for_file(entry.get("file_path"))
|
|
)
|
|
if sub_type:
|
|
entry["sub_type"] = sub_type
|
|
return entry
|
|
|
|
def get_model_roots(self) -> List[str]:
|
|
"""Get checkpoint root directories"""
|
|
return config.base_models_roots
|