feat: add checkpoint scanner integration to recipe scanner

- Add CheckpointScanner dependency to RecipeScanner singleton
- Implement checkpoint enrichment in recipe data processing
- Add _enrich_checkpoint_entry method to enhance checkpoint metadata
- Update recipe formatting to include checkpoint information
- Extend get_instance, __new__, and __init__ methods to support checkpoint scanner
- Add _get_checkpoint_from_version_index method for cache lookup

This enables recipe scanner to handle checkpoint models alongside existing LoRA support, providing complete model metadata for recipes.
This commit is contained in:
Will Miao
2025-11-21 15:36:54 +08:00
parent 4eb46a8d3e
commit 1971881537
4 changed files with 396 additions and 15 deletions

View File

@@ -9,6 +9,7 @@ from .recipe_cache import RecipeCache
from .service_registry import ServiceRegistry
from .lora_scanner import LoraScanner
from .metadata_service import get_default_metadata_provider
from .checkpoint_scanner import CheckpointScanner
from .recipes.errors import RecipeNotFoundError
from ..utils.utils import calculate_recipe_fingerprint, fuzzy_match
from natsort import natsorted
@@ -23,24 +24,39 @@ class RecipeScanner:
_lock = asyncio.Lock()
@classmethod
async def get_instance(cls, lora_scanner: Optional[LoraScanner] = None):
async def get_instance(
cls,
lora_scanner: Optional[LoraScanner] = None,
checkpoint_scanner: Optional[CheckpointScanner] = None,
):
"""Get singleton instance of RecipeScanner"""
async with cls._lock:
if cls._instance is None:
if not lora_scanner:
# Get lora scanner from service registry if not provided
lora_scanner = await ServiceRegistry.get_lora_scanner()
cls._instance = cls(lora_scanner)
if not checkpoint_scanner:
checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner()
cls._instance = cls(lora_scanner, checkpoint_scanner)
return cls._instance
def __new__(cls, lora_scanner: Optional[LoraScanner] = None):
def __new__(
cls,
lora_scanner: Optional[LoraScanner] = None,
checkpoint_scanner: Optional[CheckpointScanner] = None,
):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._lora_scanner = lora_scanner
cls._instance._checkpoint_scanner = checkpoint_scanner
cls._instance._civitai_client = None # Will be lazily initialized
return cls._instance
def __init__(self, lora_scanner: Optional[LoraScanner] = None):
def __init__(
self,
lora_scanner: Optional[LoraScanner] = None,
checkpoint_scanner: Optional[CheckpointScanner] = None,
):
# Ensure initialization only happens once
if not hasattr(self, '_initialized'):
self._cache: Optional[RecipeCache] = None
@@ -51,6 +67,8 @@ class RecipeScanner:
self._resort_tasks: Set[asyncio.Task] = set()
if lora_scanner:
self._lora_scanner = lora_scanner
if checkpoint_scanner:
self._checkpoint_scanner = checkpoint_scanner
self._initialized = True
def on_library_changed(self) -> None:
@@ -422,6 +440,9 @@ class RecipeScanner:
# Update lora information with local paths and availability
await self._update_lora_information(recipe_data)
if recipe_data.get('checkpoint'):
recipe_data['checkpoint'] = self._enrich_checkpoint_entry(dict(recipe_data['checkpoint']))
# Calculate and update fingerprint if missing
if 'loras' in recipe_data and 'fingerprint' not in recipe_data:
fingerprint = calculate_recipe_fingerprint(recipe_data['loras'])
@@ -585,6 +606,27 @@ class RecipeScanner:
return version_index.get(normalized_id)
def _get_checkpoint_from_version_index(self, model_version_id: Any) -> Optional[Dict[str, Any]]:
"""Fetch a cached checkpoint entry by version id."""
if not self._checkpoint_scanner:
return None
cache = getattr(self._checkpoint_scanner, "_cache", None)
if cache is None:
return None
version_index = getattr(cache, "version_index", None)
if not version_index:
return None
try:
normalized_id = int(model_version_id)
except (TypeError, ValueError):
return None
return version_index.get(normalized_id)
async def _determine_base_model(self, loras: List[Dict]) -> Optional[str]:
"""Determine the most common base model among LoRAs"""
base_models = {}
@@ -623,6 +665,57 @@ class RecipeScanner:
logger.error(f"Error getting base model for lora: {e}")
return None
def _enrich_checkpoint_entry(self, checkpoint: Dict[str, Any]) -> Dict[str, Any]:
"""Populate convenience fields for a checkpoint entry."""
if not checkpoint or not isinstance(checkpoint, dict) or not self._checkpoint_scanner:
return checkpoint
hash_value = (checkpoint.get('hash') or '').lower()
version_entry = None
model_version_id = checkpoint.get('id') or checkpoint.get('modelVersionId')
if not hash_value and model_version_id is not None:
version_entry = self._get_checkpoint_from_version_index(model_version_id)
try:
preview_url = checkpoint.get('preview_url') or checkpoint.get('thumbnailUrl')
if preview_url:
checkpoint['preview_url'] = self._normalize_preview_url(preview_url)
if hash_value:
checkpoint['inLibrary'] = self._checkpoint_scanner.has_hash(hash_value)
checkpoint['preview_url'] = self._normalize_preview_url(
checkpoint.get('preview_url')
or self._checkpoint_scanner.get_preview_url_by_hash(hash_value)
)
checkpoint['localPath'] = self._checkpoint_scanner.get_path_by_hash(hash_value)
elif version_entry:
checkpoint['inLibrary'] = True
cached_path = version_entry.get('file_path') or version_entry.get('path')
if cached_path:
checkpoint.setdefault('localPath', cached_path)
if not checkpoint.get('file_name'):
checkpoint['file_name'] = os.path.splitext(os.path.basename(cached_path))[0]
if version_entry.get('sha256') and not checkpoint.get('hash'):
checkpoint['hash'] = version_entry.get('sha256')
preview_url = self._normalize_preview_url(version_entry.get('preview_url'))
if preview_url:
checkpoint.setdefault('preview_url', preview_url)
if version_entry.get('model_type'):
checkpoint.setdefault('model_type', version_entry.get('model_type'))
else:
checkpoint.setdefault('inLibrary', False)
if checkpoint.get('preview_url'):
checkpoint['preview_url'] = self._normalize_preview_url(checkpoint['preview_url'])
except Exception as exc: # pragma: no cover - defensive logging
logger.debug("Error enriching checkpoint entry %s: %s", hash_value or model_version_id, exc)
return checkpoint
def _enrich_lora_entry(self, lora: Dict[str, Any]) -> Dict[str, Any]:
"""Populate convenience fields for a LoRA entry."""
@@ -827,6 +920,8 @@ class RecipeScanner:
for item in paginated_items:
if 'loras' in item:
item['loras'] = [self._enrich_lora_entry(dict(lora)) for lora in item['loras']]
if item.get('checkpoint'):
item['checkpoint'] = self._enrich_checkpoint_entry(dict(item['checkpoint']))
result = {
'items': paginated_items,
@@ -874,6 +969,8 @@ class RecipeScanner:
# Add lora metadata
if 'loras' in formatted_recipe:
formatted_recipe['loras'] = [self._enrich_lora_entry(dict(lora)) for lora in formatted_recipe['loras']]
if formatted_recipe.get('checkpoint'):
formatted_recipe['checkpoint'] = self._enrich_checkpoint_entry(dict(formatted_recipe['checkpoint']))
return formatted_recipe

View File

@@ -588,6 +588,26 @@
padding-top: 4px; /* Add padding to prevent first item from being cut off when hovered */
}
.recipe-resources-list {
display: flex;
flex-direction: column;
gap: 10px;
flex: 1;
min-height: 0;
}
.recipe-checkpoint-container {
display: flex;
flex-direction: column;
gap: var(--space-1);
}
.version-divider {
height: 1px;
background: var(--border-color);
margin: var(--space-1) 0;
}
.recipe-lora-item {
display: flex;
gap: var(--space-2);
@@ -614,6 +634,13 @@
border-left: 4px solid var(--lora-accent);
}
.recipe-lora-item.checkpoint-item {
cursor: default;
padding-top: 8px;
padding-bottom: 8px;
align-items: center;
}
.recipe-lora-item.missing-locally {
border-left: 4px solid var(--lora-error);
}
@@ -962,6 +989,10 @@
z-index: 100;
}
.badge-container .resource-action {
margin-left: auto;
}
/* Add styles for missing LoRAs download feature */
.recipe-status.missing {
position: relative;
@@ -1004,3 +1035,61 @@
.recipe-status.clickable:hover {
background-color: rgba(var(--lora-warning-rgb, 255, 165, 0), 0.2);
}
.recipe-checkpoint-meta {
display: flex;
flex-wrap: wrap;
gap: 8px;
align-items: center;
font-size: 0.85em;
margin-bottom: 2px;
}
.recipe-checkpoint-meta .checkpoint-type {
background: var(--lora-surface);
padding: 2px 8px;
border-radius: var(--border-radius-xs);
color: var(--text-color);
}
.recipe-resource-actions {
display: flex;
flex-wrap: wrap;
align-items: center;
gap: 8px;
margin-top: 2px;
}
.resource-action {
display: inline-flex;
align-items: center;
gap: 6px;
padding: 5px 10px;
border-radius: var(--border-radius-xs);
border: 1px solid var(--border-color);
background: var(--bg-color);
color: var(--text-color);
font-size: 0.9em;
cursor: pointer;
transition: background-color 0.2s ease, border-color 0.2s ease, transform 0.2s ease;
}
.resource-action.compact {
padding: 4px 10px;
font-size: 0.88em;
}
.resource-action:hover {
background: var(--lora-surface);
transform: translateY(-1px);
}
.resource-action.primary {
background: var(--lora-accent);
color: white;
border-color: var(--lora-accent);
}
.resource-action.primary:hover {
background: color-mix(in oklch, var(--lora-accent), black 10%);
}

View File

@@ -1,8 +1,11 @@
// Recipe Modal Component
import { showToast, copyToClipboard } from '../utils/uiHelpers.js';
import { showToast, copyToClipboard, sendModelPathToWorkflow } from '../utils/uiHelpers.js';
import { translate } from '../utils/i18nHelpers.js';
import { state } from '../state/index.js';
import { setSessionItem, removeSessionItem } from '../utils/storageHelpers.js';
import { updateRecipeMetadata } from '../api/recipeApi.js';
import { downloadManager } from '../managers/DownloadManager.js';
import { MODEL_TYPES } from '../api/apiConfig.js';
class RecipeModal {
constructor() {
@@ -340,6 +343,17 @@ class RecipeModal {
if (otherParamsElement) otherParamsElement.innerHTML = '<div class="no-params">No parameters available</div>';
}
const checkpointContainer = document.getElementById('recipeCheckpoint');
const resourceDivider = document.getElementById('recipeResourceDivider');
if (checkpointContainer) {
checkpointContainer.innerHTML = '';
if (recipe.checkpoint && typeof recipe.checkpoint === 'object') {
checkpointContainer.innerHTML = this.renderCheckpoint(recipe.checkpoint);
this.setupCheckpointActions(checkpointContainer, recipe.checkpoint);
}
}
// Set LoRAs list and count
const lorasListElement = document.getElementById('recipeLorasList');
const lorasCountElement = document.getElementById('recipeLorasCount');
@@ -493,6 +507,12 @@ class RecipeModal {
this.recipeLorasSyntax = '';
}
if (resourceDivider) {
const hasCheckpoint = checkpointContainer && checkpointContainer.querySelector('.recipe-lora-item');
const hasLoraItems = lorasListElement && lorasListElement.querySelector('.recipe-lora-item');
resourceDivider.style.display = hasCheckpoint && hasLoraItems ? 'block' : 'none';
}
// Show the modal
modalManager.showModal('recipeModal');
}
@@ -1047,6 +1067,177 @@ class RecipeModal {
}
}
renderCheckpoint(checkpoint) {
const existsLocally = !!checkpoint.inLibrary;
const localPath = checkpoint.localPath || '';
const previewUrl = checkpoint.preview_url || checkpoint.thumbnailUrl || '/loras_static/images/no-preview.png';
const isPreviewVideo = typeof previewUrl === 'string' && previewUrl.toLowerCase().endsWith('.mp4');
const checkpointName = checkpoint.name || checkpoint.modelName || checkpoint.file_name || 'Checkpoint';
const versionLabel = checkpoint.version || checkpoint.modelVersionName || '';
const baseModel = checkpoint.baseModel || checkpoint.base_model || '';
const modelTypeRaw = (checkpoint.model_type || checkpoint.type || 'checkpoint').toLowerCase();
const modelTypeLabel = modelTypeRaw === 'diffusion_model' ? 'Diffusion Model' : 'Checkpoint';
const previewMedia = isPreviewVideo ? `
<video class="thumbnail-video" autoplay loop muted playsinline>
<source src="${previewUrl}" type="video/mp4">
</video>
` : `<img src="${previewUrl}" alt="Checkpoint preview">`;
const badge = existsLocally ? `
<div class="local-badge">
<i class="fas fa-check"></i> In Library
<div class="local-path">${localPath}</div>
</div>
` : `
<div class="missing-badge">
<i class="fas fa-exclamation-triangle"></i> Not in Library
</div>
`;
let headerAction = '';
if (existsLocally && localPath) {
headerAction = `
<button class="resource-action primary compact checkpoint-send">
<i class="fas fa-paper-plane"></i>
<span>${translate('recipes.actions.sendCheckpoint', {}, 'Send to ComfyUI')}</span>
</button>
`;
} else if (this.canDownloadCheckpoint(checkpoint)) {
headerAction = `
<button class="resource-action primary compact checkpoint-download">
<i class="fas fa-download"></i>
<span>${translate('modals.model.versions.actions.download', {}, 'Download')}</span>
</button>
`;
}
return `
<div class="recipe-lora-item checkpoint-item ${existsLocally ? 'exists-locally' : 'missing-locally'}">
<div class="recipe-lora-thumbnail">
${previewMedia}
</div>
<div class="recipe-lora-content">
<div class="recipe-lora-header">
<h4>${checkpointName}</h4>
<div class="badge-container">${headerAction}</div>
</div>
<div class="recipe-lora-info recipe-checkpoint-meta">
${versionLabel ? `<div class="recipe-lora-version">${versionLabel}</div>` : ''}
${baseModel ? `<div class="base-model">${baseModel}</div>` : ''}
${modelTypeLabel ? `<div class="checkpoint-type">${modelTypeLabel}</div>` : ''}
</div>
</div>
</div>
`;
}
setupCheckpointActions(container, checkpoint) {
const sendBtn = container.querySelector('.checkpoint-send');
if (sendBtn) {
sendBtn.addEventListener('click', (e) => {
e.stopPropagation();
this.sendCheckpointToWorkflow(checkpoint);
});
}
const downloadBtn = container.querySelector('.checkpoint-download');
if (downloadBtn) {
downloadBtn.addEventListener('click', async (e) => {
e.stopPropagation();
await this.downloadCheckpoint(checkpoint, downloadBtn);
});
}
}
canDownloadCheckpoint(checkpoint) {
if (!checkpoint) return false;
const modelId = checkpoint.modelId || checkpoint.modelID || checkpoint.model_id;
const versionId = checkpoint.id || checkpoint.modelVersionId;
return !!(modelId && versionId);
}
async sendCheckpointToWorkflow(checkpoint) {
if (!checkpoint || !checkpoint.localPath) {
showToast('toast.recipes.missingCheckpointPath', {}, 'error');
return;
}
const modelType = (checkpoint.model_type || checkpoint.type || 'checkpoint').toLowerCase();
const isDiffusionModel = modelType === 'diffusion_model' || modelType === 'unet';
const widgetName = isDiffusionModel ? 'unet_name' : 'ckpt_name';
const actionTypeText = translate(
isDiffusionModel ? 'uiHelpers.nodeSelector.diffusionModel' : 'uiHelpers.nodeSelector.checkpoint',
{},
isDiffusionModel ? 'Diffusion Model' : 'Checkpoint'
);
const successMessage = translate(
isDiffusionModel ? 'uiHelpers.workflow.diffusionModelUpdated' : 'uiHelpers.workflow.checkpointUpdated',
{},
isDiffusionModel ? 'Diffusion model updated in workflow' : 'Checkpoint updated in workflow'
);
const failureMessage = translate(
isDiffusionModel ? 'uiHelpers.workflow.diffusionModelFailed' : 'uiHelpers.workflow.checkpointFailed',
{},
isDiffusionModel ? 'Failed to update diffusion model node' : 'Failed to update checkpoint node'
);
const missingNodesMessage = translate(
'uiHelpers.workflow.noMatchingNodes',
{},
'No compatible nodes available in the current workflow'
);
const missingTargetMessage = translate(
'uiHelpers.workflow.noTargetNodeSelected',
{},
'No target node selected'
);
await sendModelPathToWorkflow(checkpoint.localPath, {
widgetName,
collectionType: MODEL_TYPES.CHECKPOINT,
actionTypeText,
successMessage,
failureMessage,
missingNodesMessage,
missingTargetMessage,
});
}
async downloadCheckpoint(checkpoint, button) {
if (!this.canDownloadCheckpoint(checkpoint)) {
showToast('toast.recipes.missingCheckpointInfo', {}, 'error');
return;
}
const modelId = checkpoint.modelId || checkpoint.modelID || checkpoint.model_id;
const versionId = checkpoint.id || checkpoint.modelVersionId;
const versionName = checkpoint.version || checkpoint.modelVersionName || checkpoint.name || 'Checkpoint';
if (button) {
button.disabled = true;
}
try {
await downloadManager.downloadVersionWithDefaults(
MODEL_TYPES.CHECKPOINT,
modelId,
versionId,
{
versionName,
source: 'recipe-modal',
}
);
} catch (error) {
console.error('Error downloading checkpoint:', error);
showToast('toast.recipes.downloadCheckpointFailed', { message: error.message }, 'error');
} finally {
if (button) {
button.disabled = false;
}
}
}
// New method to navigate to the LoRAs page
navigateToLorasPage(specificLoraIndex = null) {
// Close the current modal

View File

@@ -57,18 +57,22 @@
<div class="info-section recipe-bottom-section">
<div class="recipe-section-header">
<h3>Resources</h3>
<div class="recipe-section-actions">
<span id="recipeLorasCount"><i class="fas fa-layer-group"></i> 0 LoRAs</span>
<button class="action-btn view-loras-btn" id="viewRecipeLorasBtn" title="View all LoRAs in this recipe">
<i class="fas fa-external-link-alt"></i>
</button>
<button class="copy-btn" id="copyRecipeSyntaxBtn" title="Copy Recipe Syntax">
<i class="fas fa-copy"></i>
</button>
</div>
<div class="recipe-section-actions">
<span id="recipeLorasCount"><i class="fas fa-layer-group"></i> 0 LoRAs</span>
<button class="action-btn view-loras-btn" id="viewRecipeLorasBtn" title="View all LoRAs in this recipe">
<i class="fas fa-external-link-alt"></i>
</button>
<button class="copy-btn" id="copyRecipeSyntaxBtn" title="Copy Recipe Syntax">
<i class="fas fa-copy"></i>
</button>
</div>
</div>
<div class="recipe-resources-list">
<div class="recipe-checkpoint-container" id="recipeCheckpoint"></div>
<div class="version-divider" id="recipeResourceDivider" style="display: none;"></div>
<div class="recipe-loras-list" id="recipeLorasList"></div>
</div>
</div>
</div>
</div>
</div>