diff --git a/py/services/recipe_scanner.py b/py/services/recipe_scanner.py index 2d11a455..1d649791 100644 --- a/py/services/recipe_scanner.py +++ b/py/services/recipe_scanner.py @@ -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 diff --git a/static/css/components/recipe-modal.css b/static/css/components/recipe-modal.css index d5e2e3fb..edcec091 100644 --- a/static/css/components/recipe-modal.css +++ b/static/css/components/recipe-modal.css @@ -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%); +} diff --git a/static/js/components/RecipeModal.js b/static/js/components/RecipeModal.js index 68d61eb9..299a2a11 100644 --- a/static/js/components/RecipeModal.js +++ b/static/js/components/RecipeModal.js @@ -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() { @@ -339,6 +342,17 @@ class RecipeModal { if (negativePromptElement) promptElement.textContent = 'No negative prompt information available'; if (otherParamsElement) otherParamsElement.innerHTML = '
No parameters available
'; } + + 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'); @@ -492,6 +506,12 @@ class RecipeModal { lorasListElement.innerHTML = '
No LoRAs associated with this recipe
'; 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 ? ` + + ` : `Checkpoint preview`; + + const badge = existsLocally ? ` +
+ In Library +
${localPath}
+
+ ` : ` +
+ Not in Library +
+ `; + + let headerAction = ''; + if (existsLocally && localPath) { + headerAction = ` + + `; + } else if (this.canDownloadCheckpoint(checkpoint)) { + headerAction = ` + + `; + } + + return ` +
+
+ ${previewMedia} +
+
+
+

${checkpointName}

+
${headerAction}
+
+
+ ${versionLabel ? `
${versionLabel}
` : ''} + ${baseModel ? `
${baseModel}
` : ''} + ${modelTypeLabel ? `
${modelTypeLabel}
` : ''} +
+
+
+ `; + } + + 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 @@ -1107,4 +1298,4 @@ class RecipeModal { } } -export { RecipeModal }; \ No newline at end of file +export { RecipeModal }; diff --git a/templates/components/recipe_modal.html b/templates/components/recipe_modal.html index fb1b74cf..bd15a9ec 100644 --- a/templates/components/recipe_modal.html +++ b/templates/components/recipe_modal.html @@ -57,18 +57,22 @@

Resources

-
- 0 LoRAs - - -
+
+ 0 LoRAs + +
+
+
+
+
+