Compare commits

..

253 Commits

Author SHA1 Message Date
Will Miao 823f71f269 feat(nodes): make Lora Stack Combiner inputs dynamic 2026-08-02 22:04:40 +08:00
Will Miao 042dd4088d fix(nodes): make Lora Stack Combiner inputs optional 2026-08-01 17:14:00 +08:00
willmiao eaa791a9eb docs: auto-update supporters list in README 2026-07-31 13:25:56 +00:00
Will Miao 2228627ff4 chore(release): bump version to v1.2.0 2026-07-31 21:25:38 +08:00
Will Miao 4c647ad9c8 fix(update): throttle nightly update badge to once per day 2026-07-31 21:18:58 +08:00
Will Miao 8ca3e6c33f fix(ui): guard marquee bulk-mode entry against click jitter and stale drag state 2026-07-31 18:40:14 +08:00
Will Miao dd6bdbf297 fix(update): persist update_channel via settings.json instead of hasGit
After b464fdc3 (preserve .git on release switch), the hasGit-based
channel detection is unreliable — .git now exists for both release
and nightly installs, so page refresh always reset the channel.

- Add _resolveChannelFromSettings() with migration heuristic:
  !hasGit → release (ZIP), detached HEAD → release (on tag),
  on branch → nightly. Uses gitInfo.branch from check-updates.
- Persist resolved channel to settings.json on first load
  (one-time migration) and on explicit switchChannel.
- Add update_channel validation (release|nightly) in backend
  update_settings handler.
- Remove hasGit-based guessing from initialize(); defer to
  checkForUpdates where full gitInfo is available.
- Channel resolution runs before checkForUpdates early-returns
  to avoid null channelMode on reload-within-interval.

Tests: 361 passed.
2026-07-31 13:23:54 +08:00
Will Miao b47dde87e4 fix(settings): suppress error toasts when optional model roots are empty 2026-07-31 10:07:52 +08:00
Will Miao 99e65cccd8 fix(update): downgrade settings backup/restore logs from INFO to DEBUG 2026-07-30 20:32:00 +08:00
Will Miao 3bdacb8f46 fix(test): update release channel git test to mock _perform_git_update instead of _download_and_replace_zip 2026-07-30 18:35:43 +08:00
Will Miao b4f9c224d3 fix(example-images): move multi→single-library consolidation to startup, eliminate per-request os.listdir()
Move reverse-migration logic from get_model_folder() (hot path, called on
every metadata/example-images request) to ExampleImagesMigration, where it
runs once at startup.  On network storage this was causing 22-38s delays
per LoRA card click.

Additionally optimize prune_stale_example_images() to read the directory
listing once instead of per image entry (O(N*M) → O(M)).  Also reorder
consolidation checks so regex filters run before filesystem stat calls.
2026-07-30 18:11:40 +08:00
Will Miao 5ec0399c81 fix(i18n): remove redundant 'preserved' sentence from release channel message, sync all 10 locales 2026-07-30 16:38:04 +08:00
Will Miao b464fdc333 fix(update): preserve .git on release channel switch, use git checkout tag
Previously, switching to the release channel would delete .git/ and
fall back to a ZIP download. This broke update.bat, manual git
commands, and CM git-based update detection.

Now the release path uses git checkout <latest-tag> when .git exists,
and only falls back to ZIP when .git is absent (CM CNR installs).
.git is never deleted - the ZIP→nightly path remains a one-way
upgrade via _init_git_repo.

Also updates locale strings (en, zh-CN, zh-TW, ja) to remove the
now-inaccurate "remove the Git repository" wording.
2026-07-29 21:23:39 +08:00
Will Miao 53825500db fix(update): add staging protection to switch_channel
switch_channel has three destructive code paths (git reset + clean,
git init + checkout --force, and rmtree + ZIP replace) that were
missing the _stage_preserved_items / _restore_preserved_items safety
net already applied to perform_update.

Wrap the channel-specific logic in a try/finally so preserved user
data (settings.json, civitai/, cache/, etc.) is physically moved
outside plugin_root before any git operation and always restored.
2026-07-29 20:41:36 +08:00
Will Miao f2ac790752 fix(update): stage preserved items outside repo before git/ZIP update
Move settings.json, civitai/, wildcards/, backups/, stats/, logs/,
cache/, and model_cache/ to a temp directory before git reset/clean
or ZIP replacement, then restore them in a try/finally block.

This prevents data loss on Windows where git clean -e exclusion
patterns can fail due to path-separator mismatches or where file
locks (open SQLite/log handles) cause the restore step to be skipped
on failure.

Also unifies three hardcoded skip lists (_clean_plugin_folder,
skip_items, skip_tracked) to derive from the single _PRESERVE_DIRS
constant, fixing drift where logs/ was missing from the ZIP path.
2026-07-29 19:49:50 +08:00
Will Miao 0d8805cdee fix(recipes): update cards in-place after LoRA download, preventing scroll reset 2026-07-29 11:35:28 +08:00
pixelpaws 656e24ac9b Merge pull request #1044 from d1udiu/fix-filter
fix(filters): prevent search query from being persisted in localStorage
2026-07-29 11:30:40 +08:00
d1udiu 6718b37403 fix(filters): prevent search query from being persisted in localStorage 2026-07-29 10:12:42 +08:00
Will Miao c9e5e784fc fix(metadata-overwrite): use sentinel default for clip_skip to accept wired 0 2026-07-28 23:13:00 +08:00
Will Miao f92f958682 fix(SaveImageLM): correct scheduler mapping and deduplicate sampler map
- Fix incorrect mapping: "normal" -> "Normal" (was "Simple")
- Replace inline sampler_mapping with CIVITAI_SAMPLER_MAP reference
  to eliminate duplicate definition
2026-07-28 21:39:09 +08:00
Will Miao f63fab0676 fix(cache): deduplicate model entries on add and reconcile to prevent duplicate cards (#1041) 2026-07-28 20:44:57 +08:00
Will Miao cfc4903c0c fix(update): read ahead_by from GitHub compare API when status is ahead/diverged
The compare API URL format compare/{local_hash}...main returns
status='ahead' when main is ahead of the local commit. The count is
in the ahead_by field, not behind_by. The old code only read behind_by
which is always 0 in this case, causing the UI to show 'Up to date'
when actually several commits behind.

Also handle status='diverged' (both sides have unique commits) by
reading ahead_by for the remote-ahead count.

Frontend adds a hash comparison fallback: if behind_by is 0 but local
and remote commit hashes differ, show 'Behind main' instead of the
incorrect 'Up to date'.

Tests: _AheadCompareDownloader and _DivergedCompareDownloader mocks
for the two status paths.
2026-07-28 17:47:38 +08:00
Will Miao a527a847fe fix(download): route UNet/diffusion model downloads to unet roots in location step
When downloading a diffusion model (UNet) from the checkpoints page, the
download modal's location step always showed checkpoint roots and paths.
Now the modal detects the file subtype and switches to unet_roots endpoint,
default_unet_root key, and 'unet' path template.
2026-07-28 17:21:12 +08:00
Will Miao 91b0bf8933 fix(download_queue): deduplicate download_history rows before creating unique index (#1041) 2026-07-27 21:36:58 +08:00
Will Miao 66d1c96783 feat(update): add Release/Nightly channel switching
- Add POST /api/lm/switch-channel endpoint with git init / ZIP fallback
- Add _backup_git/_restore_git helpers with safe rollback
- Version-info endpoint now returns has_git flag for auto-detection
- Check-updates always returns releases (changelog) regardless of channel
- Nightly mode shows 'N commits behind main' with commit hash and date
- View on GitHub link points to /commits/main in nightly mode
- Channel toggle UI with pill-style buttons in update modal
- Confirmation dialog with Esc / backdrop-dismiss support
- Channel derived from has_git on every page load, no localStorage
- i18n: 11 new keys translated across 9 non-English locales
- CSS: unified card-style sections in _base.css
- Tests: 8 new tests covering switch-channel, nightly response, init_git_repo
2026-07-27 20:27:05 +08:00
Will Miao 986128076e fix(widget): guard setValue against non-array input to prevent workflow load crash (#1039) 2026-07-26 21:54:37 +08:00
Will Miao 1de0a53241 feat(grouping): version-group library cards by HuggingFace repo for non-Civitai sources (#1040) 2026-07-26 21:46:55 +08:00
Will Miao 0ec7eaf606 fix(wildcards): resolve weighted N::value syntax inside wildcard YAML lists (#1039) 2026-07-26 18:33:31 +08:00
Will Miao d9fcb0e92b fix(filter): preserve search term through filter apply/clear operations 2026-07-26 16:49:57 +08:00
Will Miao f49b4ba4db fix(metadata-overwrite): rename 'checkpoint' input to 'model' 2026-07-26 10:59:38 +08:00
Will Miao 84e708328b fix: correct return_types propagation to GenericNodeExtractor
Two bugs prevented type-signature-based fallback from working:

- metadata_hook.py used getattr(obj.__class__, 'RETURN_TYPES')
  which fails when _async_map_node_over_list is called with
  a class (not instance) — obj.__class__ is the metaclass
  'type', which has no RETURN_TYPES. Fixed: getattr(obj, ...).

- metadata_registry.py used type(extractor) is GenericNodeExtractor
  to dispatch return_types. NODE_EXTRACTORS stores class
  references, not instances; type(Class) is always 'type',
  never the class. Fixed: extractor is GenericNodeExtractor.
2026-07-26 10:34:20 +08:00
Will Miao 125bed3f09 feat: add Metadata Overwrite node for manual generation params override 2026-07-26 08:50:21 +08:00
Will Miao 077e70169d feat: add type-signature-based fallback for unregistered nodes
GenericNodeExtractor (previously a no-op) now inspects
RETURN_TYPES to detect MODEL loaders and CONDITIONING
encoders in nodes not registered in NODE_EXTRACTORS.

- Propagate return_types from the hook layer through the
  registry to GenericNodeExtractor.extract() and update().
- MODEL detection: scan ckpt_name/unet_name/model_path/
  model_name/gguf_name fields, validate by extension.
- CONDITIONING detection: scan text/clip_l/t5xxl/prompt
  fields, store prompt text and conditioning tensor.
- _fill_missing_metadata also checks node_cache, so
  GenericNodeExtractor-handled nodes survive cache.
2026-07-25 22:14:51 +08:00
Will Miao e6dc169a05 feat: add meta hints user marks for metadata heuristic override
Users can now right-click nodes and assign meta hints
(primary_model, primary_sampler, positive_prompt,
negative_prompt) to override the metadata processor's
heuristic inference.

- Store extra_data from the API request so workflow node
  properties (including lm_marker_role) are accessible
  during metadata processing.
- _get_user_marks scans extra_data.extra_pnginfo.workflow
  for meta_* marks, falling back to prompt.original_prompt.
- extract_generation_params checks user marks before
  heuristic inference for sampler, model, and prompts.
- Warn on duplicate marks or invalid marked nodes.
2026-07-25 22:13:52 +08:00
Will Miao f34c02756d fix(recipes): eliminate O(n) fuzzy search fallback over 42k+ recipes
Drop the SequenceMatcher-based fuzzy_match fallback that froze the server
when FTS returned empty results. FTS now returns empty set for zero results
(no fallback), and when the index is not yet ready, search returns empty
rather than scanning all items in Python.
2026-07-25 17:34:37 +08:00
Will Miao 1e4c315481 fix(ModelModal): respect civitai_host setting for creator profile link 2026-07-25 07:15:12 +08:00
Will Miao a8283a0d00 fix(SaveImageLM): clarify embed_workflow tooltip — explains drag-and-drop workflow restoration
The previous tooltip was misleading: users thought workflow embedding was
automatic. New wording explains this opt-in flag stores the complete
workflow inside images, allowing one-click restoration via drag-and-drop.
PNG and WebP only.
2026-07-24 19:53:59 +08:00
Will Miao 55896669fc feat(SaveImageLM): expose webp_method and jpeg_subsampling as conditional node inputs
Add two new optional parameters to the Save Image node:

- webp_method (INT, 0-6, default 6): Controls WebP compression level.
  0=fastest/largest, 6=slowest/smallest. Previously hardcoded to 0.
- jpeg_subsampling (INT, 0-2, default 0): Controls JPEG chroma
  subsampling. 0=4:4:4 (best quality), 1=4:2:2, 2=4:2:0.

Frontend JS extension hides/disables each parameter when the
selected file_format doesn't apply (e.g., webp_method is hidden
when saving as PNG or JPEG). 7 new tests cover parameter plumbing
and default consistency across INPUT_TYPES, save_images(), and
process_image().
2026-07-24 19:32:51 +08:00
Will Miao e341e0b9d2 fix(test): update parameters assertion to include Version: ComfyUI after metadata format upgrade 2026-07-24 18:29:07 +08:00
Will Miao e6538c83bb fix(metadata): restore sha256 after hydrate_model_data to prevent KeyError in CivitAI fetch
hydrate_model_data replaces model_data with .metadata.json content which
may lack sha256 (corrupted file, concurrent write, etc.). Restore the
cached sha256 after hydration and persist the fix back to disk so
subsequent lookups don't hit the same error.

Also improve error log to include file_path for debugging.
2026-07-24 12:07:18 +08:00
Will Miao 92e1285ea5 feat(SaveImageLM): upgrade metadata output to A1111/Civitai-compatible format
- Replace plain-text Lora hashes with Hashes JSON dict matching A1111 convention
- Add Civitai resources JSON array with AIR URNs for direct model version linking
- Add Clip skip, Version: ComfyUI fields to generation params line
- Build AIR strings from local scanner cache (no API calls needed)
- Add complete sampler name mapping (CIVITAI_SAMPLER_MAP) and base model → AIR slug mapping (BASE_MODEL_AIR_SLUG) sourced from civitai ecosystem constants
- Remove lora text prepending from prompt line; LoRA info now in structured JSON sections
2026-07-24 06:20:28 +08:00
Will Miao 2aabd1d90e fix(ai): use json_schema instead of json_object for broader provider compatibility (#1033)
LM Studio and some other OpenAI-compatible servers reject
response_format=json_object but accept json_schema. Switch to the
equivalent json_schema format and add a fallback that retries
without response_format when the provider rejects the format type.
2026-07-23 09:17:29 +08:00
Will Miao 7b8b778f83 fix(widget): restore strength drag on lora entries and header
widget.value is a getter/setter that returns a new array on every read,
so handleStrengthDrag with updateWidget=false mutated a discarded copy.
Introduce __dragActive flag to suppress renderLoras in setValue during
drag, allowing mutations to persist through widget.value without
destroying the DOM. Use try-finally to guarantee flag cleanup.
2026-07-23 08:31:34 +08:00
Will Miao 7c8dc57d55 fix(security): use abspath instead of realpath in containment checks to support symlinks (#1028) 2026-07-23 07:06:41 +08:00
Will Miao fe95fae5f2 fix(workflow): include Create Hook LoRA in lora_code_update handler 2026-07-22 11:40:56 +08:00
Will Miao ce8a95abf7 chore(release): bump version to v1.1.9 2026-07-21 22:22:39 +08:00
Will Miao c8e7e543d6 fix(api): remove overstrict model type validation in getApiEndpoints
The validation in getApiEndpoints threw for page types not in
MODEL_TYPES (e.g. 'recipes'), crashing the recipes page initialization
when FilterManager calls it via createBaseModelTags(). The throw was
synchronous and outside the fetch().catch() chain, causing an uncaught
promise rejection that aborted the entire app initialization.

getApiEndpoints is a URL builder -- validation belongs to callers that
need strict type checking (they already use isValidModelType()). For
non-model-type pages like recipes, the generated URLs are correct
(the backend does have /api/lm/recipes/* routes).

Fixes regression from f53f859a (feat(filter): add debounced tag search).
2026-07-21 22:09:30 +08:00
Will Miao a9dbb15ffa fix(create_hook_lora): lazy import comfy.hooks/comfy.utils to fix CI pipeline (#744) 2026-07-21 18:44:04 +08:00
Will Miao cf64043f7d fix(security): add library root containment check for delete/move/rename operations (#1028) 2026-07-21 15:23:38 +08:00
Will Miao ccaff92c18 fix(nodes): register Create Hook LoRA node in workflow target registries 2026-07-21 14:56:39 +08:00
Will Miao 585b5c922a feat(nodes): add Create Hook LoRA (LoraManager) node for multi-LoRA hook pipelines 2026-07-21 09:44:28 +08:00
Will Miao ea80c2224c fix(download): prevent path traversal in download template resolution (#1028) 2026-07-20 21:08:43 +08:00
willmiao 8b0f56c1a6 docs: auto-update supporters list in README 2026-07-20 12:42:20 +00:00
Will Miao 8022d12f03 chore(release): bump version to v1.1.8 2026-07-20 20:42:04 +08:00
Will Miao 3939f7f91b chore: Update lora manager basic example workflow 2026-07-20 20:20:35 +08:00
Will Miao aebf2e37dd fix(filter): apply preset on full tile click and suppress i18n double-translate warnings
- Move preset apply handler from span.preset-name to div.filter-preset so
  clicking anywhere on the tile triggers the preset, not just the label text.
- Add whitespace heuristic in showToast() to skip translate() for plain
  messages that are already translated at the call site. This prevents
  i18next from logging 'Translation key not found' for pre-translated
  strings like 'Preset "name" applied'.
2026-07-20 17:57:14 +08:00
Will Miao f53f859a71 feat(filter): add debounced tag search with backend search-tags endpoint 2026-07-20 17:37:47 +08:00
Will Miao d916375abe fix(checkpoint): populate hash index from pre-computed metadata to prevent repeated hash re-calculation (#1002) 2026-07-20 12:24:54 +08:00
Will Miao 57983df4bd fix(recipe): resolve recipe metadata update bugs in cache sort, allowed fields, and bulk API routing
- Use safe .get() in RecipeCache._resort_locked instead of itemgetter to prevent KeyError when recipe missing created_date; align sort key with _sort_cache_sync (prefer modified, fallback created_date, fallback 0)
- Add base_model to allowed_fields in persistence_service.update_recipe() so the field passes validation
- Route bulk base model updates through updateRecipeMetadata() on recipes page instead of generic saveModelMetadata(), matching existing isRecipesPage pattern used in setBulkFavorites and saveBulkTags
2026-07-20 11:20:06 +08:00
Will Miao c68d7559a0 fix(widget): correct reorder drop indicator position when container is scrolled
The drop indicator top position was calculated using only
getBoundingClientRect() offsets (post-CSS-transform viewport space)
without accounting for container.scrollTop (pre-transform layout space).
This caused the indicator to drift upward as the user scrolled down,
eventually disappearing entirely.

Fixed by adding container.scrollTop to the position calculation and
only dividing the GBCR visual-diff portion by scale, since scrollTop
is already in pre-transform coordinate space.
2026-07-19 22:40:07 +08:00
Will Miao 9a8f5bf2d6 fix(ui): reposition download settings before AI provider section 2026-07-19 17:57:02 +08:00
Will Miao a8d742b031 feat(metadata): add CivArchive API toggle and provider fallback order settings
- Add enable_civarchive_api toggle (default on) to allow disabling
  CivArchive to avoid its rate-limit windows entirely
- Add metadata_provider_order dropdown with two presets:
  CivitAI → CivArchive → Archive DB (default) and
  CivitAI → Archive DB → CivArchive
- Wire both settings through backend (metadata_service, settings_manager,
  misc_handlers) and frontend (SettingsManager, state, settings modal)
- Reorder Metadata section in settings modal: toggles → status/management
  → fallback order, for natural top-down workflow
- Make update_metadata_providers() log the effective provider chain
  using actually-registered providers rather than settings assumptions
- Add 5 test cases covering all provider-combination paths
- Complete i18n translations for 6 new keys across all 9 non-English locales
2026-07-19 17:51:50 +08:00
Will Miao c27e4d1bfc feat(cache): opportunistic cache sync on metadata read with in-place update
- Add PersistentModelCache.update_single_model() for lightweight targeted
  SQL update (single row + incremental tag/hash deltas, no full table scan)
- Add ModelScanner.sync_cache_from_metadata() with compare-first logic:
  skips entirely when cache is already in sync; when stale, updates the
  entry in-place (O(1) instead of O(n) remove+append), incrementally
  adjusts tag counts/hash index/version index, and resorts only when
  sort-relevant fields changed
- Wire sync_cache_from_metadata() into BaseModelService.get_model_metadata()
  via fire-and-forget asyncio.create_task — disk I/O is already paid for
- Include identity re-validation guard against concurrent cache replacement
- Add 16 tests covering _cache_entries_differ, sync_cache_from_metadata
  (no-change, in-place, fallback, conditional resort), and
  update_single_model (insert, tag delta, hash delta)
2026-07-19 08:32:21 +08:00
Will Miao d15a8aa9a2 fix(workflow): accept non-string widget values and support GlobalSeed node in gen-params (#1026) 2026-07-18 22:57:20 +08:00
Will Miao 74a7d12ca4 fix(test): add missing options mock in LoraInfoWidget test 2026-07-18 22:14:03 +08:00
Will Miao 2f94a9773e feat(workflow): redirect gen-params updates to connected Primitive nodes (#1026)
When a KSampler marked as 'Send Gen Params Target' has widget inputs
wired to Primitive nodes (PrimitiveNode, PrimitiveInt, PrimitiveFloat,
etc.), sending gen params from the Lora Manager UI now updates the
Primitive node's value instead of the KSampler widget. This is
necessary because ComfyUI's execution engine reads from the connected
input, ignoring the widget value when a wire is present.

Also fix two minor issues found during review:
- Remove unnecessary String() wrapping on numeric gen params (seed,
  steps, cfg) to preserve native types through the JSON/WS path
- Correct misleading isNodeEnabled comment: LGraphEventMode values
  are 0=Always, 2=Never, 4=Bypass (not 'Normal/Enabled')
2026-07-18 22:10:48 +08:00
Will Miao 37bdfa21ea fix(standalone): ensure sys.path includes script dir for python_embeded compatibility (#1025) 2026-07-18 21:22:25 +08:00
Will Miao f0bf2728c9 fix(downloads): accept download_id in history delete/retry endpoints, add unique index 2026-07-18 21:10:42 +08:00
Will Miao dc715aa273 fix(download): fallback to downloadUrl when all mirrors are deleted
When Civitai returns 404 for /models/{id} (e.g. due to Civitai API bug
where un-deleted models still get 404), the fallback to CivArchive
provides metadata.  However CivArchive may return mirrors with every
entry marked deletedAt, while the file's downloadUrl is still valid.

Before this fix, _build_download_urls_from_file_info used an if/else
that skipped the downloadUrl fallback whenever the mirrors array was
non-empty, even when all mirrors were filtered out.  Now downloadUrl
is always tried when no usable mirror remains.

Also deduplicated the inline mirror-processing code at the second call
site by replacing it with a call to the shared helper.
2026-07-18 18:28:32 +08:00
Will Miao 7ee2361e87 fix(config): remove stale 'default' library entry and consolidate example images on startup 2026-07-18 17:25:36 +08:00
Will Miao e04c22f83f fix(widgets): allow text selection in LoraInfoWidget description tab 2026-07-17 18:33:59 +08:00
Will Miao 681cc13e90 fix(widgets): persist LoRA entry selection and active tab across save/load 2026-07-17 18:27:34 +08:00
Will Miao 090e0297d4 fix(downloader): hold session lock in retry paths to prevent session close race
Refactor _create_session() to make-before-break: snapshot old session,
assign new one first, then close old.  Previously, concurrent download
retries called _create_session() without the session lock (violating its
docstring contract) and closed the old session while other coroutines
held active references — causing aiohttp to raise "NoneType has no
attribute connect" when dereferencing the torn-down connector.

Also wrap the two _create_session() calls in the integrity-retry and
network-retry paths with self._session_lock to match the locking
discipline used by the session property and refresh_session().
2026-07-17 17:21:05 +08:00
Will Miao 6f71335be4 feat(widgets): add Description tab to LoraInfoWidget with dual-mode rendering support
- Add Notes/Description tab switching with tab state persistence in widget value
- Lazy-load model description and version description from /lm/loras/metadata
- Render CivitAI HTML descriptions inline via v-html
- Auto-fetch description when LoRA selection changes while on Description tab
- Fix Vue mode height containment via contain:layout size (lm-vue-node class)
- Fix scroll wheel isolation: widget scroll vs canvas zoom in both render modes
- Add docs/comfyui-dual-mode-widgets.md with widget rendering patterns
2026-07-17 15:04:47 +08:00
Will Miao 7f51812c1e feat(nodes): add LoRA Syntax → Path node (#1015) 2026-07-16 19:57:29 +08:00
Will Miao a9dc4d7b9d fix(widgets): reuse orphaned DOM containers after undo/redo in Vue render mode
In ComfyUI Vue render mode, WidgetDOM.vue reuses its component instance
during undo/redo without re-calling mountWidgetElement(), leaving newly
created widget containers detached from the DOM.

- AutocompleteTextWidget: scan for empty containers by ID prefix and reuse
- Loras widget: scan for empty .lm-loras-container elements and reuse
- Prevent duplicate event listeners by guarding listener setup on new
  containers only
- Keep container in DOM on cleanup (clearChildren instead of remove)
  so it can be found and reused by the next factory invocation
2026-07-16 18:54:00 +08:00
Will Miao 5d50ddb5d4 fix(ui): exit bulk mode after send-to-workflow completes 2026-07-16 18:54:00 +08:00
Will Miao f86198d234 fix(loras): include folder prefix in context menu and bulk send-to-workflow
When using full path lora syntax, the context menu (single/bulk)
and bulk copy actions were passing only the file basename to
buildLoraSyntax(), ignoring the folder prefix. This caused the
output to look like legacy A1111 format even when full path mode
was enabled.

Aligns all entry points with ModelCard.handleSendToWorkflow(),
which correctly includes the folder prefix.

Also fixes selectAllVisibleModels() to cache the folder field,
preventing missing prefix on select-all-then-send flows.
2026-07-16 18:54:00 +08:00
Will Miao ffe65d983c feat(api): add GET endpoints for update-lora-code and update-node-widget
Add GET variants of the two POST endpoints used by the send-to-workflow
feature. Parameters are read from query string instead of JSON body,
supporting both simple repeated node_id params and JSON-encoded node_ids
for complex graph references.
2026-07-15 21:49:22 +08:00
Will Miao b0b5be913c fix(downloads): reject re-insertion of download_ids already in history
In add_to_queue, check download_history before INSERT OR IGNORE.  Without
this check, a fire-and-forget /queue/complete failure on the extension side
would allow the same download_id to be re-inserted after complete_download()
deleted it from the queue — creating phantom queued entries for already-
finished downloads.
2026-07-15 19:12:22 +08:00
Will Miao 01efcbc584 fix(loras): allow toggle deselect on LoRA entry click 2026-07-14 18:23:06 +08:00
Will Miao 02c249917a fix(recipe): ensure custom recipes_path is added to preview allowed roots on startup 2026-07-14 18:15:20 +08:00
Will Miao 419bbc90b2 feat(lora-info): add Lora Info display node
Add a pure frontend node that shows filename and editable notes for
a selected LoRA. Connect any output from a LoRA Loader/Stacker/Randomizer/
WanVideoSelect to the lora_source input — selecting a LoRA in the source
widget updates the info display automatically.

- Python node (LoraInfoLM): display-only, no workflow execution
- Vue widget: filename label, auto-sizing notes textarea, save button
  with ComfyUI toast feedback on save
- Frontend extension: wire-based selection propagation with stale-response
  race guard; clears display on wire disconnect
- Backend: get-notes endpoint now returns file_path alongside notes;
  matching supports full-path lora syntax; fix NoneType crash in
  trigger words endpoint; document cache file_name invariant
- Wired into all four lora widget nodes (Loader, Stacker, Randomizer,
  WanVideoSelect)
2026-07-14 18:00:31 +08:00
willmiao b0c4510fdb docs: auto-update supporters list in README 2026-07-13 14:18:36 +00:00
Will Miao bf6a614e0d chore(release): bump version to v1.1.7 2026-07-13 22:18:16 +08:00
Will Miao feab01cd9c fix(preview): hide license icons for models without CivitAI metadata 2026-07-13 19:49:10 +08:00
Will Miao 966024e534 fix(registry): force re-registration on WS refresh to prevent timeout, demote empty-registry log to debug
- workflow_registry.js: add force param to refreshRegistry(), bypass fingerprint
  dedup when responding to lora_registry_refresh WS message. Without this, the
  backend's wait_for_all() times out after 0.5s because the frontend skips the
  register-nodes POST when the workflow fingerprint hasn't changed (common after
  ComfyUI restart with an empty or unchanged workflow).
- misc_handlers.py: demote 'No nodes registered after refresh' from WARNING to
  DEBUG — empty workflows are a normal operational state, not a warning-worthy
  condition.
2026-07-13 19:10:48 +08:00
Will Miao 2018722cc8 fix(registry): handle compound subgraph node IDs, add proactive node push from graph hooks
- Handle compound node IDs (e.g. "252:0") from expanded group subgraphs
  to fix 400 Bad Request on workflows with group nodes
- Frontend proactively pushes node data via afterConfigureGraph and
  LiteGraph hooks (onNodeAdded/onNodeRemoved/graphChanged), eliminating
  WebSocket round-trip latency for most "Send to Workflow" operations
- Add content-fingerprint dedup to skip duplicate register-nodes POSTs
- Fast-path cache returns immediately when tabs are registered (including
  0-node registrations), avoiding unnecessary WS refresh cycles
- Distinguish "Empty Registry" from other errors in standalone UI toast
- Reduce WS refresh timeout 2s→0.5s, add cooldown and lock to prevent
  concurrent refresh storms
- All [LM:Registry] logs at DEBUG level
2026-07-13 18:02:26 +08:00
Will Miao 9d85c2a44a fix(ui): prevent tags widget from auto-resizing in Vue mode when tags change 2026-07-13 14:55:40 +08:00
Will Miao 03dd047e62 fix(download): return 200 instead of 500 when user cancels download 2026-07-13 11:47:48 +08:00
Will Miao 86b547c1e0 fix(locales): add missing downloadStopped key to toast.downloads section 2026-07-13 11:35:48 +08:00
Will Miao bab9752c8b fix(download): close modal before progress overlay and fix downloadId ReferenceError on cancel 2026-07-13 11:29:47 +08:00
Will Miao 774cc1be86 fix(download): use file ID for exact match, add debug logging for multi-file selection (#1023)
- Frontend: send file.id in file_params, use null instead of hardcoded defaults
- Backend: priority matching (ID exact → primary → lenient metadata)
- Lenient metadata: only compare fields present on both sides (fixes GGUF size mismatch)
- Add debug logs at key points: entry, file_params received, match result, anomaly signals
2026-07-13 11:15:03 +08:00
Will Miao 234b73c8a2 feat(ui): add cancel button to download progress modal 2026-07-13 09:40:53 +08:00
Will Miao abd06c48f4 fix(settings): reject checkpoints↔unet path overlap in extra folder paths with inline error UI
Changes:
- Backend: _validate_folder_paths() now checks checkpoints↔unet overlap
  within the same library using os.path.realpath() for symlink resolution
- Backend: set() calls _validate_folder_paths() for both folder_paths and
  extra_folder_paths before writing
- Backend: extracted _normalize_path_set() helper to eliminate duplicated
  normalization logic
- Frontend: inline error display with red border + error message below the
  conflicting input, no save triggered
- Frontend: path normalization (strip trailing slash, lowercase) in pre-check
  to reduce false negatives vs backend realpath
- Frontend: asymmetric error UX — message only on the user-edited side,
  red border on the pre-existing conflict side
- CSS: has-error styles with hardcoded rgba fallback for older browsers
- i18n: checkpointUnetOverlap + checkpointUnetOverlapInline keys added to
  all 10 locale files
2026-07-13 08:22:40 +08:00
Will Miao 6ca411e4e4 fix(ui): make loras widget fixed-size with user-controlled node resize
Remove dynamic height calculation that auto-resized the node when
LoRAs are added or removed. The widget now stays at the size the user
sets via the node resize handle, scrolling when content overflows.

- Drop updateWidgetHeight() and hardcoded entry-count height math
- Set --comfy-widget-min-height once (200px) instead of recalculating
- In Vue mode: add contain:layout+size to break the ResizeObserver
  feedback loop that forced node growth with content (CSS via
  .lm-loras-container.lm-vue-node scoped to vueNodesMode only)
- Remove unused "Node 2.0: Maximum visible LoRA entries" setting
2026-07-12 22:35:58 +08:00
Will Miao 6470021e77 feat(settings): persist LORA_MANAGER_PORTABLE to settings.json on first use (#1018) 2026-07-12 09:32:30 +08:00
Will Miao 71658ab37b feat(settings): add LORA_MANAGER_PORTABLE env var for per-instance settings isolation (#1018) 2026-07-12 07:44:31 +08:00
Will Miao 4f016a8024 feat(fetch): skip CivArchive API for HuggingFace-sourced models
- Bulk refresh filter now excludes models with hf_url
- Individual refresh for HF models only checks CivitAI API
- CivArchive client validates model IDs before querying
2026-07-11 20:29:54 +08:00
Will Miao f362ed585b fix(preview): gracefully handle deleted preview files - image fallback, cache cleanup, quieter logs
- Add onerror handler on <img> previews to fallback to no-preview.png
- Fire async cache cleanup when preview file returns 404
- Add ModelCache.clear_preview_by_path() for safe stale-url removal
- Downgrade /api/lm/previews 404 log from warning to debug
2026-07-10 21:25:07 +08:00
Will Miao 196172624f fix(ui): allow autocomplete textarea resize in app mode (#1020) 2026-07-09 11:59:09 +08:00
Will Miao 316702b7ab fix(hf): allow subdirectory paths in HF resolve URLs, strip repo-internal dirs on save (#1019) 2026-07-09 09:18:38 +08:00
Will Miao a7625b009f fix(ui): also exit bulk mode after enrich-hf-llm-bulk completes 2026-07-07 20:31:16 +08:00
Will Miao 5d4a33c90d fix(hf): stop using realpath for download path construction, match CivitAI approach 2026-07-07 20:24:47 +08:00
Will Miao 041a6b8525 Revert "fix(hf): pass computed folder to _save_hf_metadata instead of re-deriving from paths"
This reverts commit 54b44131b6.
2026-07-07 20:13:20 +08:00
Will Miao 2638109ad6 feat(hf): add Link to HuggingFace feature with unified Link Model submenu
- Merge Relink to Civitai and new Link to HuggingFace into a single
  'Link Model' submenu with sub-options for each source
- Add POST /api/lm/set-hf-url endpoint to associate a model with a
  HuggingFace repo URL, saving hf_url to .metadata.json
- Add link_hf_modal.html for URL input, following relink-civitai pattern
- Use update_single_model_cache instead of add_model_to_cache to
  prevent duplicate cache entries after linking
- Remove os.path.realpath usage for consistency with relink-civitai
- Raise errors instead of silently falling back to LoRA scanner when
  model root cannot be determined
- Scope .input-group CSS rules to modal IDs to fix style conflicts
  with download-modal.css
- Add i18n keys across all 10 locales with translations for
  zh-CN, zh-TW, ja, ko, de, es, fr, he, ru
2026-07-07 20:04:47 +08:00
Will Miao b019326747 feat(ui): auto-exit bulk mode after all bulk operations complete 2026-07-06 18:51:33 +08:00
Will Miao 54b44131b6 fix(hf): pass computed folder to _save_hf_metadata instead of re-deriving from paths 2026-07-06 17:34:43 +08:00
Will Miao a1d948025c fix(hf): strip empty trainedWords from metadata JSON to keep sidecar clean 2026-07-06 16:49:51 +08:00
Will Miao a90b2514ba feat(ui): group HF batch files by repo with collapse/expand, fix nested scroll & collapse animation
- Group HF batch download files by repo with collapsible group headers
- Fix nested scrollbar conflict (inner scrollbar undraggable) by making batch-preview-list flex-fill
- Fix collapse animation glitch (items disappearing before container shrinks) by keeping expanded during max-height transition
- Visual polish: hover lift, backdrop-filter glass, design token alignment
- Remove redundant database icon from group header
- Guard transitionend handlers against rapid-click races
2026-07-06 16:36:26 +08:00
pixelpaws cb4ad27813 Merge pull request #1013 from willmiao/agent
Hugging Face model metadata AI enrichment
2026-07-06 12:21:19 +08:00
Will Miao 637831248b fix(agent): route WS error events through onError instead of dead onComplete branch 2026-07-06 12:18:17 +08:00
Will Miao 00228deaaa fix(download): retry on Civitai 429 rate limit instead of removing images from metadata
When Civitai returns 429 (Too Many Requests) during example image
downloads, the previous behavior treated all failures identically and
permanently removed the corresponding images from model metadata —
making them impossible to retry.

This commit adds:
- 429 detection + Retry-After header parsing in download_to_memory
- Exponential backoff retry (up to 3 attempts) in
  download_model_images_with_tracking
- Separate tracking of rate-limited vs permanently failed URLs
- rate_limited_models progress tracking persisted to disk
- Rate-limited models are NOT added to failed_models/processed_models
  so they are automatically retried on subsequent download runs
- Force mode clears failed_models when rate-limited images exist
2026-07-06 11:58:19 +08:00
Will Miao 2373edf73c feat(ui): load provider model catalog asynchronously to avoid blocking page render 2026-07-06 10:02:09 +08:00
Will Miao e0e1b804a7 fix(llm): require api_base for custom provider without preset default 2026-07-06 10:02:04 +08:00
Will Miao fecbe8241f fix(agent): use status= instead of status_code in json_response calls 2026-07-06 10:02:00 +08:00
Will Miao 5983eaa1ce refactor(llm): use catalog-based max_tokens, remove JSON retry, reduce Ollama num_ctx
- Parse limit.output from model catalog alongside model IDs
  for per-model max output token limits
- Use catalog lookup in chat_completion_json() to set max_tokens;
  fall back to 4096 for unknown models (e.g. local Ollama)
- Remove the JSON retry (response_format → plain text fallback);
  keep _try_salvage_json as last-resort for truncated responses
- Reduce Ollama num_ctx from 32768 to 8192 (sufficient for
  metadata enrichment, saves VRAM)
- Fix stale test comment referencing removed retry
2026-07-06 09:13:42 +08:00
Will Miao 07fa454f72 chore(tests): stop tracking HF enrichment baseline snapshots
Remove tests/enrich_hf_validation/baselines/ from git tracking
(.gitignore entry + git rm --cached). These contain README snapshots
from community HF repos that may include NSFW/sensitive content.

Local files are preserved on disk for offline reference.
2026-07-06 01:08:25 +08:00
Will Miao 4b5aa45379 chore(tests): update bash code block tests to match preserved-bash behavior
Commit 9a0d866b changed _strip_fenced_code_blocks to preserve bash/shell
code blocks (they carry CLI setup and trigger-word metadata signal).
Update the two affected tests to expect bash content in the output
instead of asserting it is stripped.

- Rename test_bash_code_block_stripped → test_bash_code_block_preserved
- Update assertions: expect 'pip install' in result
2026-07-06 01:02:04 +08:00
Will Miao 9a0d866be4 fix(agent): preserve bash/shell code blocks in readme_processor during README cleaning 2026-07-06 00:40:35 +08:00
Will Miao 308d8f71b8 feat(ui): gray out enrich-hf-llm when no hf_url, add backend fast-fail, rename labels across locales, reposition menu item 2026-07-06 00:34:18 +08:00
Will Miao d0e8938039 fix(agent): call _format_base_models via self. to prevent NameError
The bare call  inside _build_prompt_context
would raise NameError because class methods don't close over class-level
scope. Use  instead to trigger attribute lookup.

Update enrich_hf_metadata prompt.md clue locations for better LLM accuracy.
Update baseline report to v2 (mean 69.0, 46 models, +2.2pp vs baseline 71.1%).
Consolidate README snapshots into baselines/readmes/.
2026-07-06 00:10:30 +08:00
Will Miao 13ed898b6b chore(tests): add base_model ground truth mapping for all 46 test entries 2026-07-05 20:47:30 +08:00
Will Miao e1dfd1c2a6 chore(tests): add two test entries and their HF README snapshots 2026-07-05 20:45:01 +08:00
Will Miao e3e944911b refactor(agent): extract shared scanner iteration into _find_model_entry
_Previous_ _find_scanner_for_model and identify_model_type contained ~25 lines
of identical scanner-iteration + path-matching logic.  Factor it into
_find_model_entry() so a new scanner type or edge-case fix can't drift apart.
2026-07-05 18:03:57 +08:00
Will Miao 51c0135250 refactor(agent): rename agent_cli to metadata_ops, strip temp debug logs
- Rename py/agent_cli/ -> py/metadata_ops/ (module was never agent-related)
- Rename tests/agent_cli/ -> tests/metadata_ops/
- Remove 9 low-value/debug INFO log points across agent_handlers.py,
  agent_service.py, llm_service.py, and metadata_ops/__init__.py
- Keep LLM raw response at DEBUG level for diagnostics
- Consolidate per-model progress + LLM result into single concise
  log line with basename instead of full path
- Update package/class/method docstrings to clarify this is a
  pipeline infrastructure, not a true agent loop
2026-07-05 18:00:58 +08:00
Will Miao 7b19bbb14e fix(agent): preserve preview URLs for collection repo models with flat heading structure
Three-part fix for enrich_hf_metadata failing to extract correct preview_url
from HuggingFace collection repos where models share flat heading levels:

1. _strip_standalone_images() now converts <img> tags to markdown image
   syntax ![alt](src) instead of stripping the URL entirely, so the LLM
   can still extract preview URLs.

2. _extract_section() uses a line-count-based forward window (stopping at
   <a id> anchors) for non-heading matches, instead of stopping at the
   very next heading. This prevents same-level sub-headings (# Download,
   # Trigger, # Sample prompt within a single model section) from
   truncating the window before sample images are included.

3. Post-processor preview fallback now filters gallery images to the
   model-specific README section before falling back to the repo-wide
   first image.
2026-07-05 17:05:47 +08:00
Will Miao 5494a70f40 chore(tests): commit validation dataset and baseline reports into repo
Move the HF model list from ~/Documents/ into tests/enrich_hf_validation/test_data/
and commit the pipeline validation baseline artifacts (report.json,
preprocessing_audit.json, README snapshots) into baselines/.

Update config.py and run_validation.py defaults to use repo-relative paths
via os.path.dirname(__file__) instead of ~/Documents/ hardcode.

Originates from changes in 8fb00998 (validation pipeline audit).
2026-07-05 17:03:45 +08:00
Will Miao 26c9ade1c9 feat(agent): optimize base model prompt — grouped display, comprehensive mapping rules, filename inference
- agent_service._format_base_models: output bullet list instead of
  JSON array for cleaner LLM parsing
- prompt.md mapping section: replace 14-row HF→CivitAI table with
  compact rule set covering 14 mapping paths including new entries
  for HiDream-ai, OnomaAIResearch/Illustrious, ideogram-ai/ideogram,
  Tongyi-MAI/Z-Image-Turbo, and Wan-AI/Wan2.*
- base_model extraction instruction: add guidance to infer from
  model filename, YAML tags, and README body text when YAML
  frontmatter has no explicit base_model:
2026-07-05 15:45:17 +08:00
Will Miao 87db23825f feat(constants): add 12 new CivitAI base models from API, sync JS/Python abbreviations and categories 2026-07-05 11:44:53 +08:00
Will Miao 8fb00998a7 feat(agent): fix extract_relevant_section false positives, add validation pipeline audit
- extract_relevant_section: raise token threshold >3, verify anchor
  sections contain basename, require 2+ heading token overlaps, skip
  TOC-style headings (markdown links), verify heading section size
- metadata_constructor: parse repo_id,model_name.safetensors format
  so model_path basename matches real filename
- config: replace hardcoded SUPPORTED_BASE_MODELS with dynamic
  init_supported_base_models() using production list_base_models()
- preprocessing_auditor: new Phase 1.5 audit module — fetches each
  README, runs extract_relevant_section + clean_readme_for_llm,
  records stats and flags, saves raw READMEs for cross-reference
- run_validation: integrate audit phase, add --audit-only mode,
  add LLM config consistency check, add ComfyUI root to sys.path
- report_generator: add Preprocessing Audit and Config Warnings
  sections to both markdown and JSON reports
2026-07-05 11:18:48 +08:00
Will Miao dd3aa97d0a refactor(agent): rename md_to_html to readme_processor, fix section extraction, widget parsing, and list_base_models
- Rename md_to_html.py → readme_processor.py (file no longer just HTML conversion)
- _extract_section: include YAML frontmatter, use heading-level-aware forward
  walk (sub-headings under # are included), increase walk limit past 30 lines
- _is_heading: exclude </hN> closing tags from boundary detection
- _heading_level: new helper for heading-level-aware section matching
- css: yield 0 for heading like closing tags, was unexpectedly caught by _is_heading
- extract_gallery_images: fix YAML block scalar (text: >-) prompt extraction;
  use endswith instead of == to detect the block marker
- _strip_widget_section: add to clean_readme_for_llm (widget text is handled
  by post-processor, not needed in LLM prompt)
- _strip_standalone_images: keep markdown image URLs intact for LLM preview
  extraction (was stripping to alt text only)
- list_base_models: switch from scanner-cache aggregation to
  CivitaiBaseModelService.get_base_models() - always returns full list
- Ollama: add num_ctx=32768 to payload options so thinking models have room
  to both reason and produce output
- Add tests/agent_cli/test_readme_processor.py: 59 tests covering extraction,
  cleaning, section matching, heading detection
- Update existing tests for behavioral changes
2026-07-05 06:39:54 +08:00
Will Miao 8bee8f4069 fix(recipe): fallback to locate custom example image on disk by model hash and image id (#1012) 2026-07-04 18:40:34 +08:00
Will Miao 817fe21b3e fix(ui): read cfg_scale and clip_skip with snake_case fallback, pass custom image id for recipe creation (#1012) 2026-07-04 18:40:24 +08:00
Will Miao 905c37290f chore: update runtime logs to use 'LLM enrichment' instead of 'Agent skill'
- agent_handlers.py: 'Agent skill' -> 'LLM enrichment' in all log messages
- skill_registry.py: 'agent skills' -> 'prompt-based skills' in discovery log
- llm_service.py: docstring 'agent skills' -> 'LLM-based enrichment features'
2026-07-04 16:53:41 +08:00
Will Miao f7632a47f9 feat(agent): enrich_hf_metadata with per-model progress and in-place card update
- PostProcessor returns updates dict from enrich_hf_metadata
- AgentService includes updated_data per model in WebSocket progress events
- Convert preview_url to HTTP URL via config.get_preview_static_url()
- LoraContextMenu: showEnhancedProgress + updateSingleItem per model
- BulkContextMenu: same pattern, remove window.location.reload()
- Guard empty updated_data and clean up callbacks on HTTP error
2026-07-04 16:50:56 +08:00
Will Miao 646f1ddfb1 refactor(agent): align 'Agent' naming to 'AI/LLM' to match current implementation
- locales/en.json: 'Enrich Metadata (Agent)' -> 'Enrich Metadata (AI)'
- Rename SKILL.md -> prompt.md with backward compat in skill_registry.py
- JS context menu action IDs: enrich-hf-agent -> enrich-hf-llm
- HTML template data-action attributes synced to match
- docstring cleanup: 'agent skill' -> 'skill pipeline' / 'feature'
2026-07-04 14:06:50 +08:00
Will Miao 170c8068c5 feat(agent): enrich_hf_metadata — filename-aware section matching, preview extraction for markdown/HTML/widget, JSON salvage, instance_prompt fallback, and validation suite
- extract_relevant_section(): trim README to model-filename-matching section
  for collection repos (download link, anchor ID, heading strategies)
- _strip_standalone_images(): preserve markdown image URLs so LLM can
  extract preview_url; strip only HTML <img> tags
- extract_simple_markdown_images(): extract civitai.images from ![]() body
- extract_html_img_tags(): extract from <img src="..."> (deadman44-style)
- extract_gallery_images(): fix widget parser for YAML - output: dash prefix
- _is_heading: exclude </hN> closing tags from boundary detection
- _extract_section: start at matching heading when match IS a heading line
- _try_salvage_json(): recover truncated JSON (close braces/brackets in
  LIFO order, close unterminated strings, strip trailing commas)
- PostProcessor: store _llm_confidence, add instance_prompt YAML fallback
- agent_service: pass model_basename to prompt, trim README via
  extract_relevant_section before clean_readme_for_llm
- Add tests/enrich_hf_validation/ suite: 100-model pipeline with progress
  checkpoint/resume, per-field scoring, markdown+JSON reporting
- Fix evaluation_engine: read _llm_confidence (not _llm_response)
2026-07-04 12:00:15 +08:00
Will Miao 3494037d20 fix(download): pass proxy to aria2 for actual file transfers (#1010) 2026-07-04 11:07:18 +08:00
Will Miao a1fd4e150b feat(agent): optimize enrich_hf_metadata with README cleaning, Ollama native API, and expanded fields
- Add clean_readme_for_llm() to strip noise from README before LLM injection
- Keep widget section text (valuable tag signal) and unmarked code blocks (trigger words)
- Preserve standalone image alt text instead of removing entirely
- Switch Ollama to native /api/chat with think:false to fix empty content on thinking models
- Extract Sample Gallery table images and deduplicate with widget images
- Only strip code blocks with explicit language tags (bash)
- Add notes and usage_tips fields to SKILL.md output format and post-processor
- Clean up dead code, fix regex edge cases, remove double type annotation
2026-07-04 08:01:50 +08:00
Will Miao b22f09bd1d fix(standalone): load extra folder paths from library settings in standalone mode 2026-07-03 19:21:56 +08:00
Will Miao 4ed9169646 feat(ui): redesign AI Provider settings with provider presets and model catalog
- Replace hardcoded provider list with PROVIDER_PRESETS (OpenAI, Ollama,
  DeepSeek, Groq, OpenRouter, OpenCode Go, Custom)
- Load model lists from models.dev/api.json catalog at startup
- Add Combobox vanilla JS component for model/base-URL selection
- Fetch local Ollama models via live API instead of catalog
- Hide API key values from frontend (boolean-only llm_api_key_set)
- Add i18n translations for all 9+ locales
- Update snapshot tests for new response fields
2026-07-03 16:08:51 +08:00
Will Miao f06c60bd47 fix(agent): handle plain YAML scalar text in extract_gallery_images
Widget entries with unquoted multi-line YAML scalars (e.g. "text: two samurais...\n  continuation") were not parsed, leaving gallery image prompts empty. Add a third branch for plain scalar format alongside the existing quoted and >- folded block handlers.
2026-07-03 07:34:24 +08:00
Will Miao ee8250c26c feat(agent): extract HF widget gallery images into civitai.images with recommended dimensions
- Add extract_gallery_images() to parse YAML widget entries from README
  frontmatter, convert relative image URLs to absolute HF URLs, and
  build civitai.images-compatible entries with prompt metadata
- LLM now extracts recommended_width/recommended_height from README
  (e.g. "Best Dimensions"), used as gallery image dimensions
- extract_gallery_images() accepts default_width/height parameters,
  falling back to 512x512 when LLM provides no recommendation
- Frontend ShowcaseView.js: defensive NaN guard for 0 width/height
- post_processor: consistently merge civitai updates across triggers,
  description, and gallery blocks with distinct variable names
- SKILL.md: add recommended_width/recommended_height to output schema
- 62 tests pass, including gallery extraction and dimension tests
2026-07-03 07:07:19 +08:00
Will Miao 88349bf944 feat(agent): render HF README as HTML in modelDescription, move converter to skill-local module
- Add inline convert_readme_to_html() in new skill-local md_to_html.py
  (zero external deps, handles h1-h4/bold/italic/code/lists/tables/links/hr)
- Strip YAML frontmatter, <Gallery />, badge images, HTML comments pre-conversion
- Fix indented whitespace after lists being misidentified as code blocks
- Fix HTML double-escaping in _inline_md (each pattern escapes independently)
- LLM short_description → civitai.description ("About this version" sidebar)
- raw README HTML → modelDescription (description tab, always available offline)
- Pass full readme_content from agent_service to post_processor
- 51 tests for converter + 4 updated/added post-processor tests
2026-07-02 23:34:52 +08:00
Will Miao a8adcaf023 feat(agent): improve enrich_hf_metadata skill with priority_tags, preview_url fix, civitai.trainedWords
- Add identify_model_type() helper to determine lora/checkpoint/embedding
- Pass priority_tags from user settings to LLM prompt for tag relevance
- SKILL.md: instruct LLM to exclude technical/generic HF tags, cross-reference
  against priority_tags; forbid ['None'] placeholder for trigger words
- post_processor: fix preview_url not updated after download (now writes local
  .webp path to metadata); write trigger words to civitai.trainedWords instead
  of top-level; sanitize ['None']/'null'/'n/a' placeholder values to []
- download_preview() now returns str | None (local path) instead of bool
- Update tests for new return type and nested civitai.trainedWords structure
2026-07-02 22:14:44 +08:00
Will Miao 63785f82b5 refactor(agent): consolidate skill definition into single SKILL.md with YAML frontmatter
Merge skill.yaml (metadata) and prompt.md (prompt template) into a
single SKILL.md file with YAML frontmatter, matching the agent-skill
convention used by opencode and Claude Code.

- Add frontmatter parser (_parse_skill_file) to SkillRegistry
- Remove skill.yaml, prompt.md, empty skills/__init__.py
- Remove obsolete load_handler method
- Update tests for new format and cleaned-up fields
2026-07-02 21:29:02 +08:00
Will Miao cf898da193 feat(agent): add LLM-powered metadata enrichment system with AgentCLI and PostProcessor
Introduce an agent skill framework for LLM-driven metadata enrichment:

- AgentCLI (py/agent_cli/): in-process wrappers around internal services
  using standard relative imports, eliminating the need for sys.path hacks
- LLMService: centralized BYOK (bring-your-own-key) LLM client supporting
  OpenAI, Ollama, and custom OpenAI-compatible endpoints
- PostProcessor: deterministic engine that applies LLM output via AgentCLI
  (replaces old handler.py + _BASE_MODEL_ALIASES approach)
- SkillRegistry: filesystem-based skill discovery (skill.yaml + prompt.md)
- AgentService: orchestrates skill execution with WebSocket progress
- Frontend AgentManager: WebSocket listeners, skill execution, config UI
- Context menu entries (single + bulk) for "Enrich Metadata (Agent)"
- Settings UI for AI Provider configuration (BYOK)
- Full i18n support across 9 locales

Bug fixes found during review:
- aiohttp.web.json_response: status_code= -> status=
- settings_modal cancelEditApiKey: wrong argument position
- AgentManager.isLlmConfigured: allow Ollama without API key
- PostProcessor._merge_tags: lowercase all tags to match TagUpdateService
2026-07-02 21:27:01 +08:00
Will Miao 3c83e78d9f feat(ui): auto-newline after pasting URL in download and batch-import textareas
Extract auto-newline-on-paste logic into shared setupAutoNewlineOnPaste() utility in uiHelpers.js.
Apply it to both the Download modal (modelUrl) and Batch Import modal (batchUrlInput)
textarea, so users can paste multiple URLs in succession without manually pressing Enter.
2026-07-02 10:53:33 +08:00
Will Miao d7291f73c9 fix(download): recognize civitai.red and civitai.green URLs in batch download (#1003) 2026-07-02 10:28:03 +08:00
Will Miao fe90f7f9b1 feat(ui): add searchable base model dropdown with filename inference in model modal
Replace native <select> with a searchable dropdown that:
- Filters options as the user types
- Shows filename-inferred suggestions at the top in a "Suggested" section
- Supports keyboard navigation (ArrowUp/Down/Enter/Escape)
- Allows typing custom values not in the list
- Removes dead .base-model-selector CSS

Adds 3 new i18n keys (baseModelSearchPlaceholder, baseModelSuggested,
baseModelNoMatch) with translations for all 9 locales.
2026-07-01 14:31:08 +08:00
Will Miao 8b344ea39f feat(ui): add View on Hugging Face button, plumb hf_url through full cache pipeline 2026-07-01 08:38:16 +08:00
Will Miao 8348a0cef8 fix(download): harden HF download path validation, fix WebSocket leak, add URL detection tests (#965, #977)
Security hardening:
- Validate repo format with strict regex (reject .. traversal)
- Validate filename rejects path separators and ..
- Validate relative_path rejects absolute paths and ..
- Verify model_root is within configured scanner roots using
  realpath + os.sep guard to prevent prefix-match bypass
- Add realpath-based escape detection for final dest_path

Bug fixes:
- Fix WebSocket leak in _downloadHfSingle: wrap ws.close() in
  try/finally so it closes even if downloadHfModel() throws
- Same fix for batch HF download per-file WebSocket loop

Frontend hardening:
- Tighten HF repo regex: require huggingface.co for full URLs,
  reject bare .. patterns
- Add 12 unit tests for detectUrlType() covering HF resolve,
  HF repo, CivitAI, CivArchive, direct HTTP, edge cases
2026-07-01 05:51:58 +08:00
Will Miao 7cf785b72f fix(ui): unify HF file selection UI, remove cloud icon, add select-all, cleanup dead code (#965, #977)
- Unify single-URL and multi-URL HF repo flows to use the same batch
  preview interface (remove separate repoFileStep)
- Remove unnecessary cloud icon from HF batch preview items
- Use formatFileSize() instead of hardcoded MB text
- Change default selection to unchecked (no preselected files)
- Add select all / deselect all checkbox with dynamic Next button
- Clean up dead CSS, HTML template, and JS methods from removed
  repoFileStep
- Add selectAll i18n key with translations for all 10 locales
- Fix batch progress bar name fallback for HF items
2026-06-30 23:28:35 +08:00
Will Miao e8913f4481 feat(ui): dynamically populate base model dropdown from CivitAI API, add Krea 2 constants (#1001) 2026-06-30 22:41:17 +08:00
Will Miao f9c3d8dc97 fix(metadata): demote CivArchive hash lookup failure from ERROR to DEBUG
A model not being found on CivArchive by hash is a routine case (the
model simply isn't published there), not an error. The callers already
log the outcome at WARNING (bulk_metadata_refresh) or DEBUG
(metadata_sync_service) with full context, making this ERROR-level log
both misleading and redundant.
2026-06-30 19:42:30 +08:00
Will Miao 09ca91fc0e feat(download): add Hugging Face model download to standalone UI wizard (#965, #977)
Integrate HF model downloading into the existing CivitAI-style wizard flow:
- URL type detection (civitai / hf-resolve / hf-repo / direct-http)
- Repo file explorer with checkbox-based file selection
- Batch/queue download with per-file WebSocket progress
- Aria2 backend support (respects download_backend setting)
- Scanner cache integration via create_default_metadata + add_model_to_cache
- i18n updates for all 10 locales
2026-06-30 19:36:12 +08:00
Will Miao 16f5222efd fix(cache): prevent corrupted cache rows from breaking model listings (#730)
Cache corruption (NULL model_name/file_name from legacy DB rows or partial
writes) caused format_response to raise KeyError/AttributeError, failing the
entire /loras/list request and showing no models in the UI.

Fix across three layers:
- format_response (lora/checkpoint/embedding): replace direct dict[] access
  with .get() fallbacks; return None for entries missing file_path
- handlers: filter None entries from list/excluded/fetch/duplicate/conflict
  endpoints instead of letting them crash or appear as null in responses
- model_scanner: always use validate_batch repaired copies (previously
  discarded when no invalid entries, leaving None values in raw_data)
- persistent_model_cache: add or-empty-string guards on read and write for
  nullable TEXT columns (model_name, file_name, folder, base_model, etc.)
2026-06-30 09:02:42 +08:00
Will Miao 28e7c04b37 fix(settings): migrate all settings subdirectories on portable mode switch 2026-06-29 21:40:37 +08:00
Will Miao 28f99c46d3 fix(update): preserve user data dirs during Git-based update via git clean -e excludes
git clean -fd in _perform_git_update deleted untracked, non-ignored
directories (wildcards, stats, backups, civitai, caches, logs) during
portable-mode updates, since released tags do not list them in .gitignore.
Add -e excludes for all user-managed paths to both nightly and stable
update branches. Add regression tests for both paths.
2026-06-29 21:10:38 +08:00
Will Miao 205194f4e6 chore: add stats, wildcards, backups, and logs dirs to .gitignore 2026-06-29 19:46:04 +08:00
willmiao 402d8b07cf docs: auto-update supporters list in README 2026-06-28 14:17:19 +00:00
Will Miao 3e303ab316 chore(release): bump version to v1.1.6 2026-06-28 22:17:02 +08:00
Will Miao e9e8c31ad1 fix(registry): store nodes per-client to prevent multi-tab race condition
Move NodeRegistry from a single global _nodes dict to a per-client
(_tab_nodes) structure so that multiple ComfyUI browser tabs no
longer overwrite each other's workflow node data during a
lora_registry_refresh cycle.  The merged result is a union of all
known tabs' target nodes, eliminating the non-deterministic failure
where send-to-workflow could randomly target a tab lacking valid
targets.

- NodeRegistry.register_nodes(sid, nodes) replaces per-tab data
  without affecting other tabs.
- NodeRegistry.get_merged_registry() returns the union across all
  connected clients, together with tab_count / per-tab metadata.
- prepare_for_refresh() snapshots the current active sockets; caller
  re-reads before merging so that newly-connected tabs are not pruned.
- workflow_registry.js sends api.clientId in the POST body so the
  backend can identify which tab is registering.
2026-06-28 17:57:58 +08:00
Will Miao 703a6a4ea0 fix(import): request withMeta=true from CivitAI API, fix checkpoint type guard and CivArchive version lookup
- Add &withMeta=true to image info URL so API returns full generation
  metadata (resources with hash/type) instead of null meta
- Fix checkpoint assignment guard: check modelId instead of id so non-
  checkpoint types (upscaler) are not wrongly set as recipe checkpoint
- Skip modelVersionIds loop when resources/civitaiResources already
  provided LoRAs, preventing hash-resolved duplicates
- Fix int/str type comparison in CivArchive get_model_version so
  version ID matching works correctly
2026-06-27 22:22:48 +08:00
Will Miao 283730cf38 fix(import): discover LoRA + checkpoint from modelVersionIds when API meta is null
When CivitAI image API returns meta=null and modelVersionIds at root
level, the import flow now:

- Injects modelVersionIds + browsingLevel into a minimal metadata dict
  so the parser can discover LoRAs and checkpoints (both import-from-url
  and analyze-image paths)
- Adds checkpoint dedup + fallback in the parser's modelVersionIds
  handler to avoid duplicate API calls
- Runs EXIF extraction unconditionally in analyze-image path, then
  merges with API metadata (fixes gen params loss)
- Propagates preview_nsfw_level through all three import paths:
  import-from-url, analyze-image (UI Import), and batch-import,
  plus the frontend save flow
2026-06-27 17:05:38 +08:00
Will Miao 20417797e8 fix(download): accept UNet and Diffusion Model file types from CivitAI
- Prefer file type (UNet/Diffusion Model) over baseModel name when
  deciding whether a checkpoint routes to the unet folder
- Add UNet to backend primary file type whitelist
- Add Krea 2 to DIFFUSION_MODEL_BASE_MODELS
- Include UNet/Diffusion Model files in frontend file selection UI
- Use actual file type from CivitAI in download params instead of
  hardcoded 'Model'
2026-06-27 08:56:11 +08:00
Will Miao 004c69b9ef fix(marquee): use document coordinates, add auto-scroll, support VirtualScroller off-screen cards
- Convert marquee selection from viewport to document coordinates so
  scrolling during a drag no longer deselects off-screen cards.
- Add RAF-based auto-scroll when dragging near viewport edges.
- Compute off-screen card positions from VirtualScroller layout
  parameters instead of relying on DOM queries.
2026-06-27 08:21:21 +08:00
Will Miao 47fe2d3783 chore: remove deprecated reference files from refs/ 2026-06-27 07:02:22 +08:00
Will Miao 36ef840a22 fix(parser): merge Lora hashes over empty Hashes JSON values and skip entries without hash 2026-06-26 22:31:36 +08:00
Will Miao 09c2445ac9 fix(ui): prevent scroll jump on model card click caused by sort dropdown focus
The document-level click handler in SortDropdown.js called trigger.focus()
unconditionally on every click outside the sort group. When a model card
was clicked to open the modal, focus() triggered scrollIntoView on the
.sort-trigger button, perturbing .page-content.scrollTop and causing the
card grid to jump up a few pixels.

The same interference also broke the back-to-top smooth-scroll animation:
frame-by-frame focus/scroll perturbations caused VirtualScroller to
schedule repeated re-renders, interrupting the compositor-thread scroll.

Fix: only return focus to the trigger when the dropdown was actually open,
so ordinary page clicks (e.g. clicking a model card) never force focus.
2026-06-26 19:40:12 +08:00
Will Miao 8a6d23f9c7 Revert "fix(ui): replace smooth scroll with instant for back-to-top to avoid VirtualScroller conflict"
This reverts commit a429e6b1c3.
2026-06-26 19:36:08 +08:00
Will Miao 3d207b6744 fix(updates): mark cross-folder versions as in-library during folder-filtered refresh (#997)
When refreshing updates with a folder filter, versions already present in
other folders were excluded from the is_in_library check, making them
appear as available updates. When the user tried to download, the global
check found the file already exists and returned 'model already exists'.

Fix by also collecting the cross-folder version set when folder_path is
provided, and using the union (folder-filtered + cross-folder) for
is_in_library in both _build_record_from_remote and
_merge_with_local_versions.
2026-06-26 17:40:41 +08:00
Will Miao b3edda62ad refactor(ui): persist sort per-mode with two storage keys, add recipes sort persistence 2026-06-26 17:07:17 +08:00
Will Miao a429e6b1c3 fix(ui): replace smooth scroll with instant for back-to-top to avoid VirtualScroller conflict
The back-to-top button used scrollTo({top:0, behavior:'smooth'}) which
conflicts with VirtualScroller's DOM manipulations during the smooth
scroll animation. Each animation frame triggered handleScroll() ->
scheduleRender() -> renderItems(), causing the browser to interrupt
the smooth scroll animation mid-way, resulting in only ~1 page of
upward scroll instead of reaching the top.

Root cause: commit 311e89e9 fixed VirtualScroller to listen on the
correct scroll container (.page-content), but this meant every scroll
event during smooth animation now triggers expensive DOM operations
that abort the browser's compositor-thread smooth scroll animation.

Fix: use instant scroll (scrollTop = 0) so the position is set
immediately without triggering frame-by-frame VirtualScroller
interference.
2026-06-26 16:31:31 +08:00
Will Miao c1bf9c6221 test(aria2): verify _wait_until_ready captures stderr on subprocess early exit
Regression test for the pipe-race bug where _drain_stderr consumed
aria2's error output before _wait_until_ready could read it.
2026-06-26 14:41:32 +08:00
Will Miao 75fffc1e25 fix(aria2): move stderr drain after _wait_until_ready to avoid swallowing startup errors
_drain_stderr and _wait_until_ready both read from the same stderr pipe.
Starting the drain task before _wait_until_ready creates a race where the
drain task consumes aria2's early-exit error message before the startup
waiter can read it, resulting in an empty error message in the logs.

Also confirmed that --fsync does not exist as an aria2 option (exit code
28 = Invalid argument).
2026-06-26 14:32:43 +08:00
Will Miao f264bab65c fix(aria2): remove --fsync=false to avoid crash on older aria2c versions
Exit code 28 (Invalid argument) indicates this user's aria2c does not
support the --fsync option. Remove it unconditionally; the stderr drain,
relaxed RPC timeouts, and increased retry coverage remain in place.
2026-06-26 14:24:46 +08:00
Will Miao 154fcd803b fix(aria2): disable fsync and relax RPC timeouts to prevent aria2 freeze on large files
aria2 default --fsync=true calls fsync() after each write, which blocks
the entire single-threaded process on large files under Docker overlay.
Add --fsync=false to eliminate this blocking source.

Relax aiohttp session timeout: total=30 → sock_connect=10, sock_read=60
so that transient I/O delays don't cut off legitimate tellStatus RPCs.

Increase retry params (4 attempts, 3s delay) to give aria2 more recovery
time when blocked on synchronous I/O.
2026-06-26 14:19:37 +08:00
Will Miao 4ef32d3a96 fix(ui): prevent bulk-mode highlight from being clipped on edge cards 2026-06-26 11:59:28 +08:00
Will Miao d2d109a69c feat(ui): replace native sort select with custom dropdown sized to selected text 2026-06-26 09:53:04 +08:00
Will Miao 3a2941d751 fix(aria2): drain stderr pipe to prevent aria2 freeze, retry RPC status on transient failure
Root cause: aria2c subprocess stderr pipe (64 KB buffer) was never
drained. When enough error/warning output accumulated, aria2's write()
blocked, freezing the entire process including its RPC handler. The
tellStatus call then timed out after 30s with asyncio.TimeoutError(),
producing the empty error message in 'Failed to query aria2 download
status: '.

Fixes:
- Drain stderr in a background task so pipe never fills up
- Retry get_status() RPC calls up to 3 times on transient failure
- In the failure path, preserve .safetensors when .aria2 is absent
  (the download was likely complete on disk)
2026-06-26 08:25:05 +08:00
Will Miao 0ac10dfd42 fix(ui): prevent Launch LoRA Manager button from disappearing when opening properties panel in subgraph (#996) 2026-06-25 20:47:29 +08:00
Will Miao 9c95856b2f fix(trigger-wheel): prevent Vue render mode from intercepting strength wheel events
In Vue render mode, ComfyUI's TransformPane uses a capture-phase wheel
handler (@wheel.capture) that fires before the tag element's bubble-phase
strength-adjustment listener. It checks wheelCapturedByFocusedElement(),
which requires data-capture-wheel on a focused element. The tag divs had
data-capture-wheel but were not focusable, so the check failed, causing
the capture handler to forward the event to the canvas (triggering zoom)
and stopPropagation() which prevented the strength handler from running.

Fix: move data-capture-wheel from individual tags to the container, make
it focusable (tabIndex=-1), and add a window-level capture-phase wheel
listener that focuses the container before TransformPane checks it.
2026-06-25 14:58:20 +08:00
Will Miao 5ce4667d32 feat(node-marker): add 🎯 emoji prefix to Mark as context menu item 2026-06-24 22:36:45 +08:00
willmiao be53fda6df docs: auto-update supporters list in README 2026-06-24 14:11:36 +00:00
Will Miao f48de05102 chore(release): bump version to v1.1.5 2026-06-24 22:11:17 +08:00
Will Miao 93ad81ed87 fix(ui): replace full-page loading overlay with grid-scoped loader to eliminate flicker
- Add .grid-loading-overlay CSS: position:absolute inside card grid,
  semi-transparent dark background, z-index 100, pointer-events:none
- Add showGridLoading() / hideGridLoading() to VirtualScroller:
  creates/removes the scoped overlay inside the card grid only
- Modify loadMoreWithVirtualScroll(): replace full-page
  state.loadingManager overlay with grid-scoped loading, defer
  hide via requestAnimationFrame to eliminate blank-frame gap
- Clean up gridLoadingOverlay in dispose() to prevent DOM leak
2026-06-24 21:11:13 +08:00
Will Miao ea14d211be refactor(ui): unify search bar placeholder to i18n key header.search.placeholder
- Replace page-specific header.search.placeholders.* keys with a single
  header.search.placeholder key (value: "Search", no ellipsis)
- Keep header.search.notAvailable for the statistics page
- Remove unused placeholder/placeholders/notAvailable entries from all
  10 locale files; preserve options and searchIn keys
- Update Jinja template and JS header to use the new unified key
2026-06-24 20:30:38 +08:00
Will Miao 8052cefd46 feat(ui): add keyboard shortcut cue in search bar, fix clear button positioning 2026-06-24 20:21:15 +08:00
Will Miao 845815b9b7 fix(flash): fix text widget flash in Vue mode, add fade and hover dismissal
- Fix Vue mode: text widgets (CLIPTextEncode, Prompt LM) had no
  [data-testid=widget-layout-field-label], so findRowEl never matched.
  Added fallback strategies: bare <label> text match and widget index match.
- Fix Vue mode: flash background pulse was never applied — @keyframes was
  defined but no rule bound it to .lm-flash. Replaced with CSS transition
  on .lm-flash-host class for value text color fade-in/fade-out.
- Fix Vue mode: -webkit-text-fill-color set by ComfyUI overrode
  even with !important. Added -webkit-text-fill-color override to .lm-flash.
- Fix canvas mode: highlight rect was double-offset because onDrawForeground
  ctx is pre-translated to node.pos. Removed background rect entirely per
  design decision; kept text_color + inline color only.
- Add fade-in (250ms) / fade-out (400ms) for text color in both modes.
  Canvas-drawn widgets use rAF color interpolation; DOM widgets use CSS
  transition. Fixed hexToRgb to handle 3-digit hex shorthand (#DDD).
- Add hover dismissal to canvas mode via app.canvas.getWidgetAtCursor().
  Vue mode already had it via mouseover listener.
- Replace 60fps rAF poll with 100ms setInterval for hover detection.
- Fix batch cleanup closure bug: isDomWidget evaluated per-widget instead
  of per-call; fade rAF cancellers tracked per-widget in _lmFadeCancels map.
- Unify flash color from #66B3FF to LM brand accent #4299E0.
- Fix Vue fade-out: keep .lm-flash-host 300ms after removing .lm-flash so
  CSS transition persists. Canvas DOM widgets: keep inline transition 300ms
  after clearing color.
2026-06-24 19:35:30 +08:00
Will Miao 609dc5d783 feat(sort): enable versions_count sort in non-grouped mode
Sort by Most/Fewest versions first now works when Group by model is off.

- Backend: group items by modelId (respecting version_grouping setting),
  count versions per group, sort groups by count, expand groups with
  versions sorted by version id descending
- CSS: remove rule that hid the sort option in non-grouped mode
- Tests: add 3 tests covering desc, asc, and same_base variants
2026-06-24 17:14:39 +08:00
Will Miao 7a71b34b54 feat(vlm): sort versions by newest first in VLM view, with disabled sort dropdown
When viewing all versions of a model (VLM mode via 'x versions' button):
- Backend always sorts by version ID descending, ignoring current sort_by
- A temporary 'Newest version first' option is injected into the sort
  dropdown (removed on exit, not a permanent option)
- The sort dropdown is disabled (greyed out) while VLM is active
- On clearing VLM, the previous sort preference is restored and the
  dropdown re-enabled
- Handles stale VLM state (e.g. after page reload with leftover session)
- Covers all three model page types: loras, checkpoints, embeddings

Also fixes review nits:
- Correct i18n call pattern (defaultValue in options object)
- Shared _restoreSortAfterVlm() helper to avoid triple duplication
2026-06-24 16:25:14 +08:00
Will Miao 71a459422f feat: send gen params to workflow with visual cues
- Add genParamsMapper.js: sampler/scheduler display→internal mapping,
  combined-name parsing, widget matching
- Add sendGenParamsToWorkflow() in uiHelpers.js: resolves sampler,
  fetches registry by send_gen_params marker, sends via update-node-widget
- Add send-params-btn UI in showcase hover panel and recipe modal
- Add flashWidget() in workflow_registry.js: text-color visual cue
  on updated widget values (Vue: inline style + CSS, canvas: property shadow)
- Add silent option to sendWidgetValueToNodes for consolidated toast
- Normalize param display labels (cfg_scale→CFG, etc.) in recipe modal
- Add 33 tests for genParamsMapper; update existing test assertions
2026-06-24 15:39:57 +08:00
Will Miao cd2628a0ee feat(ui): add send-prompt-to-workflow button for prompt and negative prompt
- Add sendPromptToWorkflow() and stripLoraTags() exports to uiHelpers.js
- Add send button (paper-plane icon) to recipe modal and showcase hover panel
- Restructure showcase metadata panel layout to match recipe modal style
- Respect strip <lora:> setting before sending
- Uses 'replace' mode (not append) on text-capable workflow nodes
- Add translations for all 10 locales
2026-06-23 21:36:24 +08:00
Will Miao 85da7175bc feat: add Node Marker system with right-click marking 2026-06-23 20:54:32 +08:00
Will Miao d3bf0a164b fix(gitignore): add .reasonix/ to ignore list 2026-06-23 07:06:15 +08:00
Will Miao afb6ca1b8d refactor(settings): rename update_flag_strategy to version_grouping with migration 2026-06-22 16:59:32 +08:00
Will Miao 94f43426d7 feat(ui): show version count in group-by-model cards, add versions_count sort, no-reload VLM
- group_by_model dedup now counts versions per group and attaches
  version_count; respects update_flag_strategy (same_base) by
  sub-grouping on base_model
- Card footer shows clickable 'x versions' link instead of version
  name when grouped (hides HIGH/LOW badges); clicking triggers
  View Local Versions without page reload
- Added 'Local Versions' sort option (versions_count), auto-hidden
  when group_by_model is off
- Sort preference is saved/restored separately for normal and
  grouped modes
- VLM flow (triggerVlmView, clearCustomFilter) uses resetAndReload()
  via API instead of window.location.reload()
- Fixed cache mutation bug: version_count is now set on a shallow
  copy, not the cached dict, preventing stale version_count leaking
  into VLM responses
- i18n: all 9 locale files translated
2026-06-22 16:02:12 +08:00
Will Miao 2b361f4f5d feat(ui): add group-by-model toggle to global context menu
Adds a 'Group by Model' toggle entry to the right-click global context
menu for quick access, complementing the existing setting in
Settings → Layout Settings. The menu item shows a checkmark indicator
reflecting the current state and immediately reloads the view on toggle.

Also fixes he.json translation that was mojibake (garbled characters).

Includes:
- Context menu HTML item with check-indicator
- JS toggle logic via settingsManager
- i18n for all 10 locales
- Hebrew translation fix
2026-06-22 11:31:15 +08:00
Will Miao 7438072f8c feat(save-image): add %batch_num% support in batch loop 2026-06-22 09:11:38 +08:00
Will Miao 26c54fd358 fix(versions): scope VLM custom filter per-page to prevent cross-page leak
Store the originating page type alongside VLM data in sessionStorage;
validate it on every page load before applying the filter or showing
the indicator. Stale data is auto-cleaned on mismatch.

This prevents the 'View all local versions' custom filter from leaking
into the checkpoints (or embeddings) page, which caused an empty grid.
2026-06-21 12:02:06 +08:00
Will Miao 7cb6b04c63 chore: remove duplicate _truncateText from LorasControls/CheckpointsControls, add backend test for civitai_model_id filter 2026-06-21 11:19:54 +08:00
Will Miao fc29cde82a feat(versions): add View all local versions button to model versions tab
Clicking the button closes the modal, writes filter params to sessionStorage,
and reloads the page to show all local versions of the model as individual
cards (bypassing group-by-model dedup). The filter respects the update flag
strategy and the versions-filter-toggle state (same-base vs all versions).

Supporting changes:
- sessionStorage keys vlm_model_id / vlm_model_name / vlm_base_model
- BaseModelApiClient._addModelSpecificParams adds civitai_model_id param
- LoraApiClient calls super._addModelSpecificParams for VLM detection
- LorasControls / CheckpointsControls clearCustomFilter checks VLM first
- PageControls.checkVlmFilter shows customFilterIndicator with label
- Backend parses civitai_model_id, filters before group_by_model dedup
2026-06-21 11:13:53 +08:00
Will Miao 559ca946dc feat(models): add group-by-model option to collapse multiple versions into one card
Adds a 'Group by Model' toggle in Layout Settings. When enabled, only the
latest version (highest civitai.id) of each Civitai model is shown as a
single card — older versions sharing the same modelId are hidden.

Backend dedup runs in BaseModelService.get_paginated_data() before
filtering/pagination, ensuring correct paginated results. The setting
is persisted via the existing settings pipeline and passed as a query
parameter to the listing endpoint.

Includes:
- Backend: dedup logic, route param parsing, settings default
- Frontend: API param, SettingsManager wiring, toggle UI
- i18n: translations for all 10 locales
- Tests: unit test covering dedup on/off and standalone items
2026-06-21 08:48:42 +08:00
Will Miao 2b8e7c7504 fix(tests): update recipes page tests for unified controls template
- Inject #customFilterIndicator DOM in beforeEach (raw template
  renderer doesn't process Jinja2 {% include %} tags)
- Fix selector from #customFilterText to .customFilterText
2026-06-20 06:55:47 +08:00
Will Miao 6816d75933 refactor(recipes): unify controls and breadcrumb UI with model pages
- Replace inline controls+breadcrumb in recipes.html with shared includes
- Add page_id conditionals in controls.html to adapt buttons per page type
- Unify customFilterText selector to class-based in recipes.js
- Add [data-action="find-duplicates"] event listener for unified button
- Fix i18n keys to use recipes-specific translations on recipes page
2026-06-19 22:41:50 +08:00
willmiao b58abbad7c docs: auto-update supporters list in README 2026-06-19 10:31:18 +00:00
Will Miao 999814ca87 chore(release): bump version to v1.1.4 2026-06-19 18:31:03 +08:00
Will Miao 3c2760a803 fix(stats): sort Base Model Distribution X-axis labels alphabetically (#796) 2026-06-19 17:29:33 +08:00
Will Miao 0edbd7bcca fix(metadata): add LoraTextLoaderLM extractor so SaveImageLM records its loras (#801) 2026-06-19 17:13:48 +08:00
Will Miao 21e89fa7de fix(tags): normalize tag case on save and make filtering case-insensitive (#727)
- save_metadata_updates now trims/lowercases/dedupes tags on write
- ModelFilterSet tag matching is now case-insensitive (both include/exclude)
- Removed redundant .lower() calls in tag_update_service.py
2026-06-19 16:42:09 +08:00
Will Miao 968d6d1d1f feat(tags): unify recipe modal tag UI with model modal
- Replace recipe modal's custom tag display/edit with shared
  renderCompactTags/setupTagEditMode from ModelTags and utils
- Remove 300+ lines of duplicated tag display and editing code
- Parameterize setupTagEditMode with saveHandler/onSaved/showSuggestions
  options for recipe-specific save flow (updateRecipeMetadata + dirty state)
- Scope all DOM queries in ModelTags.js via options.container / this.closest
  to prevent cross-modal element conflicts
- Fix edit button alignment (justify-content: flex-start)
- Fix tag tooltip selector scoping in setupTagTooltip
- Add width: 100% to #recipeTagsContainer for edit container full width
2026-06-19 16:31:27 +08:00
Will Miao cf0fd0e0ad feat(i18n): internationalize dynamic insights content with key/params architecture (#489) 2026-06-19 13:49:03 +08:00
Will Miao 16e5dcf7b2 feat(i18n): internationalize statistics page strings across all locales 2026-06-19 13:37:01 +08:00
Will Miao ab6bb25d46 fix(example-images): skip hidden files in path validation, show offending items on failure (#807) 2026-06-19 11:54:55 +08:00
Will Miao 07f49559be fix(virtual-scroll): avoid full reload on move-to-folder, scroll to top on filter/page reset
- MoveManager/SidebarManager: replace resetAndReload with in-place
  VirtualScroller update after move operations (remove non-visible,
  update visible items' file_path). Preserves scroll position and
  avoids empty grid.
- VirtualScroller: add removeMultipleItemsByFilePath for efficient
  batch removal with Array.isArray guard.
- baseModelApi: scroll to top on loadMoreWithVirtualScroll(true),
  covering filter/sort/search/folder/views changes.
- SidebarManager selectFolder: scroll now handled centrally.
2026-06-19 09:18:49 +08:00
Will Miao b24b1a7e57 feat(settings): hide API key from frontend, use status+edit instead of password field
Backend changes:
- Add civitai_api_key to _NO_SYNC_KEYS, return only boolean civitai_api_key_set
- Clean up known template placeholder on load to prevent false positive

Frontend changes:
- Replace type=password with type=text + CSS masking (-webkit-text-security)
- Replace pre-filled input with status display (Configured/Not configured)
- Add inline edit view with Save/Cancel buttons
- Re-add eye toggle via CSS class toggle (not type switching)
- Use CSS transitions for smooth status/edit view switching

This prevents Chromium/Vivaldi password manager from triggering
'save password' prompts when opening the settings modal.
2026-06-19 08:05:04 +08:00
Will Miao faf64f8986 fix(css): migrate duplicates component to canonical color tokens
Replace undefined --lora-accent-l/c/h and --lora-warning-l/c/h with
canonical --color-accent-l/c/h and --color-warning-l/c/h from the
design token system. Fix 5 border-color declarations missing oklch()
wrapper, fix var() space syntax error in .group-toggle-btn:hover,
and replace hardcoded green with --color-success token.
2026-06-18 22:41:46 +08:00
Will Miao a617487a43 fix(ui): lift theme popover out of header stacking context to appear above modals 2026-06-18 22:19:36 +08:00
Will Miao 3012a7aef3 fix(settings): prevent Firefox save-password prompt from API key input
- Remove server-side value='...' from password field in settings modal template
  so the API key is never baked into the DOM at page load time
- Populate the input dynamically via loadSettingsToUI() when modal opens
- Clear both API key and proxy password fields on modal close to prevent
  Firefox from detecting pre-filled password fields on page navigation
2026-06-18 21:57:03 +08:00
Will Miao 499e19de34 fix(modals): tone down batch summary modal styling - remove icons, flatten gradients, lock to design tokens
- Metadata Fetch Summary: remove per-card icons, demote total/duration cards
  to neutral border, drop title icon, fix table header border width
- Batch Import Summary: replace 3em centered hero with inline left-aligned
  layout, flatten progress bar gradient, simplify circular badges to plain
  colored icons, unify border widths to 4px and token namespace to --color-
- Lock all off-scale em typography to --text-{xs,lg} design tokens
2026-06-18 21:56:58 +08:00
Will Miao 9161762ca9 fix(sidebar): align hidden indicator height (48px) and icon size with sidebar header 2026-06-18 21:14:35 +08:00
Will Miao 9bbd26efe6 feat(license-icons): add second set of license icons matching current CivitAI design
- Add 5 new Tabler SVG icons (currency-dollar, brush, user, git-merge, license)
- Implement Set 2 rendering in ModelModal.js (standalone UI) with green/red
  permission indicators and preview_tooltip.js (ComfyUI widget)
- Add use_new_license_icons setting (default: true) with toggle in settings UI
- ComfyUI tooltip reads setting directly from preview-url API response to
  eliminate race conditions and respect standalone settings changes
- Remove the now-unused separate ComfyUI setting loramanager.license_icon_style
- Add CSS for both standalone (lora-modal.css) and widget (lm_styles.css)
- i18n: translate licenseIcons keys into all 10 supported languages
- Fix test to use classic style explicitly for continued coverage
2026-06-18 21:07:44 +08:00
Will Miao 258b2622d5 fix(sidebar): align restore indicator with sidebar header and add first-use breathing animation (#990) 2026-06-18 19:22:38 +08:00
Will Miao 80ec9085dd fix(theme): replace Gruvbox with Midnight, fix accent/info hue collisions and hardcoded colors
- Replace Gruvbox preset with Midnight (deep blue-purple, violet accent)
- Fix accent/info hue collisions in Nord, Monokai, Dracula, Solarized
- Fix Solarized error/warning collision (error-h 25->5) and WCAG contrast
- Make --color-skip-refresh-* follow --color-warning-h dynamically
- Replace hardcoded rgba(24,144,255) in onboarding.css with --color-accent
- Replace hardcoded #00B87A in import modals with --color-success
2026-06-18 18:57:53 +08:00
Will Miao c5c7373e10 feat(theme): add 5 preset color themes (Nord/Gruvbox/Monokai/Dracula/Solarized) with popover selector
Implements Approach C (dual-attribute: data-theme + data-theme-preset),
keeping all 106 existing [data-theme="dark"] overrides unchanged.

- Colors: 5 professionally designed oklch palettes in tokens/colors.css
- UI: popover theme selector with mode (Light/Dark/Auto) + preset grid
- JS: cycleTheme(), setPreset(), localStorage persistence
- Locale: 12 new translation keys across 10 languages
- Polish: solid accent swatches matching flat token-driven aesthetic
2026-06-18 09:53:40 +08:00
Will Miao b7721866e5 fix(stats): implement Model Types chart in Collection tab with correct type distribution 2026-06-18 06:48:46 +08:00
Will Miao 8314b9bedb feat(downloads): add /downloads/queue/status endpoint and integrate queue lifecycle
- New GET /api/lm/downloads/queue/status handler for non-terminal status
  transitions (queued -> downloading, downloading -> paused, etc.)
- Queue lifecycle auto-integration in DownloadManager._download_with_semaphore:
  downloading -> SQLite update_status('downloading') on semaphore acquire
  completed -> complete_download('completed') on success
  canceled -> complete_download('canceled') on CancelledError
  failed -> complete_download('failed') on Exception
- All queue operations wrapped in try/except to never break the download flow
2026-06-17 23:04:30 +08:00
Will Miao 75298a402f chore(release): bump version to v1.1.3 2026-06-17 17:52:56 +08:00
Will Miao 92b5efd414 fix: guard posix_fadvise on non-Linux platforms to prevent AttributeError on Windows (#988) 2026-06-17 17:22:10 +08:00
Will Miao 33ee392b7b feat(settings): redesign Card Overlay Blur range slider to match settings UI style 2026-06-17 15:24:14 +08:00
Will Miao 5237f8b7dc chore: remove keyboard navigation UI elements and related code
- Delete static/css/components/keyboard-nav.css entirely
- Remove @import of keyboard-nav.css from style.css
- Remove keyboard-nav-hint divs from controls.html and recipes.html
- Clean up all keyboard.* translation keys from 10 locale files

The actual keyboard scrolling handlers (PageUp/PageDown in infiniteScroll.js
and VirtualScroller.js) are kept as they provide core scroll functionality.
2026-06-17 15:07:34 +08:00
Will Miao 5107313fd1 revert: restore &logo=github parameter to release-date badge
This reverts commit 95bbc669efb1aa0c23b94be6f0a5e7a188f1c019.

The real issue was shields.io GitHub API token pool exhaustion (intermittent),
not the &logo=github parameter. All 3 badges (Discord, Release, Release Date)
were affected at various times due to the same root cause: shields.io
temporarily unable to query GitHub API.
2026-06-17 11:24:40 +08:00
Will Miao 95bbc66919 fix: remove broken logo parameter from release-date badge URL 2026-06-17 11:21:26 +08:00
Will Miao e268e59419 chore: stop tracking .docs/ and add to .gitignore
.docs/ is now excluded from git tracking so working/research notes
can live there without being committed.
2026-06-17 11:20:19 +08:00
willmiao 547e1f9498 docs: auto-update supporters list in README 2026-06-17 01:57:52 +00:00
Will Miao bf32d8b6fd chore(release): bump version to v1.1.2 2026-06-17 09:57:37 +08:00
Will Miao 8299881024 refactor(sidebar): remove pin/unpin and global hide, use per-page hide only
- Remove pin/unpin and auto-hide hover mechanism (isPinned, isHovering,
  hoverTimeout, showSidebar/hideSidebar, updateAutoHideState, etc.)
- Remove global show_folder_sidebar setting (SettingsManager,
  PageControls, recipes, backend default)
- Simplify sidebar visibility to a single per-page toggle:
  · Dedicated chevron-left button in header to hide sidebar
  · Edge indicator (chevron-right) to restore when hidden
  · No dropdown, no hover area, no pin button
- Add _migrateOldSettings() to convert old sidebarPinned and
  show_folder_sidebar states to per-page sidebarDisabled
- Fix sidebar flicker on page load: CSS defaults to off-screen,
  JS explicitly sets .visible or .hidden-by-setting
- Remove obsolete CSS classes: auto-hide, hover-active, collapsed
- Remove i18n keys: pinSidebar, unpinSidebar, moreOptions
- Update test mocks for the new initialize() interface
2026-06-17 09:49:24 +08:00
Will Miao da02268196 fix(css): add top margin to stat-cards container for consistent spacing 2026-06-17 08:24:03 +08:00
Will Miao 8c4b9a1e70 fix(metadata-sync): persist not-found flags to SQLite cache on deleted-provider path
When a model is already classified as civitai_deleted=True via
.metadata.json but re-enters the failure block through the
civarchive/sqlite provider path (not the default provider),
needs_save was never set to True because civitai_api_not_found
and sqlite_attempted were both False. The flags were never
persisted to SQLite, causing the model to be re-fetched on
every restart.

Also demoted duplicate INFO/ERROR logging in fetch_and_update_model
to DEBUG (the use case already logs at WARNING), and added
exc_info=True to the fetch_all_civitai error handler.
2026-06-17 08:22:24 +08:00
Will Miao 0906c484e9 fix: actually halt bulk operations on cancel — frontend AbortController + backend guards (#986) 2026-06-17 07:20:32 +08:00
Will Miao 4199c30fec fix(metadata-sync): downgrade "Model not found" to INFO and replace model_name with file+sha256 in log 2026-06-17 00:06:43 +08:00
Will Miao 4a8084cdbc feat(save-image): support %NodeTitle.WidgetName% placeholders and fix %seed% None fallback (#314) 2026-06-16 23:48:44 +08:00
Will Miao 6263e6848c fix: move posix_fadvise(DONTNEED) after read loop so it actually evicts pages (#985) 2026-06-16 23:12:02 +08:00
Will Miao 58c266ad07 fix(scanner): respect lazy hash for checkpoints, add posix_fadvise, cancel on shutdown (#985) 2026-06-16 23:00:23 +08:00
Will Miao 2939813e1a feat(metadata-fetch): add result summary modal with i18n, fix contrast and counting bugs (#38) 2026-06-16 22:38:50 +08:00
Will Miao a9e5ee7e79 fix: follow-up nits for AVIF/JXL brotli support
- Fix JXL container ftyp size check (==20 → >=16) to accept
  wider range of valid JXL files
- Add brotli decompression size limit (2 MB) to prevent OOM
- Add trailing newline to requirements.txt
- Add unit tests for new ISOBMFF/brotli extraction paths:
  JXL/AVIF happy paths, missing brob, corrupt payload,
  non-ISOBMFF fallthrough, write-skip on AVIF/JXL,
  JSON dict/list fields, and oversized decompression
2026-06-16 16:27:56 +08:00
Will Miao a17b0e9901 Merge pull request #982 from koloved/main
Add AVIF and JXL image support with brotli metadata decompression
2026-06-16 16:24:30 +08:00
s.ivanov 8f23d966bf Update requirements.txt 2026-06-16 07:27:32 +02:00
Will Miao 7a76fc72d0 fix(rate-limit): continue to next provider on CivArchive 429 to prevent bulk refresh from freezing (#983)
When CivArchive returns HTTP 429 with a large retry_after, the bulk
metadata refresh would block for hours because:

1. FallbackMetadataProvider raised RateLimitError instead of continuing
   to the next provider (e.g., SQLite archive was never reached).

2. _RateLimitRetryHelper retried long-rate-limit 429s 3 times — all
   futile since the hourly cap hasn't reset.

3. The batch loop had no awareness of persistent rate-limiting,
   causing 192+ models to each hammer the same rate-limited endpoint.

Changes:
- FallbackMetadataProvider: all 6 methods now continue to next provider
  on RateLimitError instead of raising (model_metadata_provider.py)
- fetch_and_update_model: deleted-model path also continues on
  RateLimitError so sqlite provider gets a chance (metadata_sync_service.py)
- _RateLimitRetryHelper: when retry_after >= 120s, only 1 attempt is
  made — retries are futile for hour-scale rate limits
- BulkMetadataRefreshUseCase: tracks consecutive rate-limit failures
  and aborts early after 3 (bulk_metadata_refresh_use_case.py)

Tests: updated test_fallback_respects_retry_limit for new continue
behavior; added tests for large/small retry_after thresholds.
2026-06-16 13:08:34 +08:00
Will Miao 518a4dd5ee chore: add reasonix.toml and .codegraph/ to .gitignore 2026-06-16 13:05:11 +08:00
s.ivanov 2b6d4e5d8b Add AVIF and JXL image support with brotli metadata decompression 2026-06-15 09:28:49 +02:00
302 changed files with 54920 additions and 25327 deletions
-153
View File
@@ -1,153 +0,0 @@
# Recipe Batch Import Feature Design
## Overview
Enable users to import multiple images as recipes in a single operation, rather than processing them individually. This feature addresses the need for efficient bulk recipe creation from existing image collections.
## Architecture
```
┌─────────────────────────────────────────────────────────────────┐
│ Frontend │
├─────────────────────────────────────────────────────────────────┤
│ BatchImportManager.js │
│ ├── InputCollector (收集URL列表/目录路径) │
│ ├── ConcurrencyController (自适应并发控制) │
│ ├── ProgressTracker (进度追踪) │
│ └── ResultAggregator (结果汇总) │
├─────────────────────────────────────────────────────────────────┤
│ batch_import_modal.html │
│ └── 批量导入UI组件 │
├─────────────────────────────────────────────────────────────────┤
│ batch_import_progress.css │
│ └── 进度显示样式 │
└─────────────────────────────────────────────────────────────────┘
┌─────────────────────────────────────────────────────────────────┐
│ Backend │
├─────────────────────────────────────────────────────────────────┤
│ py/routes/handlers/recipe_handlers.py │
│ ├── start_batch_import() - 启动批量导入 │
│ ├── get_batch_import_progress() - 查询进度 │
│ └── cancel_batch_import() - 取消导入 │
├─────────────────────────────────────────────────────────────────┤
│ py/services/batch_import_service.py │
│ ├── 自适应并发执行 │
│ ├── 结果汇总 │
│ └── WebSocket进度广播 │
└─────────────────────────────────────────────────────────────────┘
```
## API Endpoints
| 端点 | 方法 | 说明 |
|------|------|------|
| `/api/lm/recipes/batch-import/start` | POST | 启动批量导入,返回 operation_id |
| `/api/lm/recipes/batch-import/progress` | GET | 查询进度状态 |
| `/api/lm/recipes/batch-import/cancel` | POST | 取消导入 |
## Backend Implementation Details
### BatchImportService
Location: `py/services/batch_import_service.py`
Key classes:
- `BatchImportItem`: Dataclass for individual import item
- `BatchImportProgress`: Dataclass for tracking progress
- `BatchImportService`: Main service class
Features:
- Adaptive concurrency control (adjusts based on success/failure rate)
- WebSocket progress broadcasting
- Graceful error handling (individual failures don't stop the batch)
- Result aggregation
### WebSocket Message Format
```json
{
"type": "batch_import_progress",
"operation_id": "xxx",
"total": 50,
"completed": 23,
"success": 21,
"failed": 2,
"skipped": 0,
"current_item": "image_024.png",
"status": "running"
}
```
### Input Types
1. **URL List**: Array of URLs (http/https)
2. **Local Paths**: Array of local file paths
3. **Directory**: Path to directory with optional recursive flag
### Error Handling
- Invalid URLs/paths: Skip and record error
- Download failures: Record error, continue
- Metadata extraction failures: Mark as "no metadata"
- Duplicate detection: Option to skip duplicates
## Frontend Implementation Details (TODO)
### UI Components
1. **BatchImportModal**: Main modal with tabs for URLs/Directory input
2. **ProgressDisplay**: Real-time progress bar and status
3. **ResultsSummary**: Final results with success/failure breakdown
### Adaptive Concurrency Controller
```javascript
class AdaptiveConcurrencyController {
constructor(options = {}) {
this.minConcurrency = options.minConcurrency || 1;
this.maxConcurrency = options.maxConcurrency || 5;
this.currentConcurrency = options.initialConcurrency || 3;
}
adjustConcurrency(taskDuration, success) {
if (success && taskDuration < 1000 && this.currentConcurrency < this.maxConcurrency) {
this.currentConcurrency = Math.min(this.currentConcurrency + 1, this.maxConcurrency);
}
if (!success || taskDuration > 10000) {
this.currentConcurrency = Math.max(this.currentConcurrency - 1, this.minConcurrency);
}
return this.currentConcurrency;
}
}
```
## File Structure
```
Backend (implemented):
├── py/services/batch_import_service.py # 后端服务
├── py/routes/handlers/batch_import_handler.py # API处理器 (added to recipe_handlers.py)
├── tests/services/test_batch_import_service.py # 单元测试
└── tests/routes/test_batch_import_routes.py # API集成测试
Frontend (TODO):
├── static/js/managers/BatchImportManager.js # 主管理器
├── static/js/managers/batch/ # 子模块
│ ├── ConcurrencyController.js # 并发控制
│ ├── ProgressTracker.js # 进度追踪
│ └── ResultAggregator.js # 结果汇总
├── static/css/components/batch-import-modal.css # 样式
└── templates/components/batch_import_modal.html # Modal模板
```
## Implementation Status
- [x] Backend BatchImportService
- [x] Backend API handlers
- [x] WebSocket progress broadcasting
- [x] Unit tests
- [x] Integration tests
- [ ] Frontend BatchImportManager
- [ ] Frontend UI components
- [ ] E2E tests
+15 -1
View File
@@ -7,17 +7,24 @@ py/run_test.py
.vscode/
cache/
civitai/
stats/
wildcards/
backups/
logs/
node_modules/
coverage/
.coverage
model_cache/
# agent
# agent / dev tooling
.opencode/
.claude/
.sisyphus/
.codex
.omo
reasonix.toml
.reasonix/
.codegraph/
# Vue widgets development cache (but keep build output)
vue-widgets/node_modules/
@@ -26,3 +33,10 @@ vue-widgets/dist/
# Hypothesis test cache
.hypothesis/
# Working/research notes (not committed)
.docs/
# HF enrichment validation baseline snapshots (contain potentially
# NSFW README content fetched from community model repos)
tests/enrich_hf_validation/baselines/
+8 -1
View File
@@ -102,6 +102,7 @@ npm run test:coverage # Generate coverage report
- ComfyUI: `app.registerExtension()`, `node.addDOMWidget(name, type, element, options)`
- Event handlers via `addEventListener` or widget callbacks
- Shared utilities: `web/comfyui/utils.js`
- Dual-mode rendering patterns (canvas vs Vue): see `docs/comfyui-dual-mode-widgets.md`
### Vue Composables Pattern
@@ -136,7 +137,13 @@ npm run test:coverage # Generate coverage report
- Dual mode: ComfyUI plugin (folder_paths) vs standalone (settings.json)
- Detection: `os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1"`
- Run `python scripts/sync_translation_keys.py` after adding UI strings to `locales/en.json`
- Symlinks require normalized paths
- Symlinks require normalized paths.
**Business paths vs real paths**: All stored paths and operation routing use the
original paths as they appear under configured model roots — symlinks are NOT
resolved. `os.path.realpath` is only for scanner dedup and the symlink cache.
Any path passed to `os.remove`/`os.rename`/`shutil.move` or validated by a
containment check MUST use the business path (i.e. `os.path.abspath`, not
`realpath`).
## Git / Commit Messages
+33 -3
View File
File diff suppressed because one or more lines are too long
+18
View File
@@ -15,6 +15,10 @@ try: # pragma: no cover - import fallback for pytest collection
from .py.nodes.lora_pool import LoraPoolLM
from .py.nodes.lora_randomizer import LoraRandomizerLM
from .py.nodes.lora_cycler import LoraCyclerLM
from .py.nodes.lora_info import LoraInfoLM
from .py.nodes.lora_syntax_to_path import LoraSyntaxToPath
from .py.nodes.create_hook_lora import CreateHookLoraLM
from .py.nodes.metadata_overwrite import MetadataOverwriteLM
from .py.metadata_collector import init as init_metadata_collector
except (
ImportError
@@ -56,6 +60,16 @@ except (
"py.nodes.lora_randomizer"
).LoraRandomizerLM
LoraCyclerLM = importlib.import_module("py.nodes.lora_cycler").LoraCyclerLM
LoraInfoLM = importlib.import_module("py.nodes.lora_info").LoraInfoLM
LoraSyntaxToPath = importlib.import_module(
"py.nodes.lora_syntax_to_path"
).LoraSyntaxToPath
CreateHookLoraLM = importlib.import_module(
"py.nodes.create_hook_lora"
).CreateHookLoraLM
MetadataOverwriteLM = importlib.import_module(
"py.nodes.metadata_overwrite"
).MetadataOverwriteLM
init_metadata_collector = importlib.import_module("py.metadata_collector").init
NODE_CLASS_MAPPINGS = {
@@ -75,6 +89,10 @@ NODE_CLASS_MAPPINGS = {
LoraPoolLM.NAME: LoraPoolLM,
LoraRandomizerLM.NAME: LoraRandomizerLM,
LoraCyclerLM.NAME: LoraCyclerLM,
LoraInfoLM.NAME: LoraInfoLM,
LoraSyntaxToPath.NAME: LoraSyntaxToPath,
CreateHookLoraLM.NAME: CreateHookLoraLM,
MetadataOverwriteLM.NAME: MetadataOverwriteLM,
}
WEB_DIRECTORY = "./web/comfyui"
+497 -416
View File
File diff suppressed because it is too large Load Diff
+208
View File
@@ -0,0 +1,208 @@
# Agent Skills System
The LoRA Manager agent skills system enables LLM-powered metadata enrichment and other AI-driven tasks. Users configure their own LLM provider (BYOK), and skills are executed through right-click context menu actions.
## Architecture
```
┌──────────────────────────────────────────────┐
│ LoRA Manager Backend │
│ │
│ ┌──────────────┐ ┌────────────────┐ │
│ │ LLMService │───▶│ LLM Provider │ │
│ │ (BYOK config, │◀───│ (OpenAI/Ollama │ │
│ │ API calls) │ │ /custom) │ │
│ └───────┬───────┘ └────────────────┘ │
│ │ │
│ ┌───────▼───────────────────────┐ │
│ │ AgentService │ │
│ │ (orchestration: validate │ │
│ │ → LLM call → post-process │ │
│ │ → WebSocket broadcast) │ │
│ └───────┬───────────────────────┘ │
│ │ │
│ ┌───────▼───────────────────────┐ │
│ │ SkillRegistry │ │
│ │ ┌─────────────────────────┐ │ │
│ │ │ enrich_hf_metadata: │ │ │
│ │ │ - skill.yaml │ │ │
│ │ │ - prompt.md │ │ │
│ │ │ - handler.py │ │ │
│ │ └─────────────────────────┘ │ │
│ └───────────────────────────────┘ │
└──────────────────────────────────────────────┘
```
### Key Design Principle
**Skills define *what* to do (prompt + post-processing). The AgentService handles *how* (LLM calls, validation, progress).**
Skills never call the LLM directly. This keeps BYOK configuration centralized and provider-agnostic.
## BYOK Configuration
Users configure their LLM provider in **Settings → AI Provider**:
| Setting | Description | Example |
|---|---|---|
| `llm_provider` | Provider type | `openai`, `ollama`, or `custom` |
| `llm_api_key` | API key (not needed for local Ollama) | `sk-...` |
| `llm_api_base` | Custom API base URL (empty = provider default) | `https://api.openai.com/v1` |
| `llm_model` | Model name | `gpt-4o-mini` |
Environment variable overrides: `LLM_API_KEY`, `LLM_MODEL`, `LLM_API_BASE`, `LLM_PROVIDER`.
### Supported Providers
- **OpenAI**: Uses `https://api.openai.com/v1` by default
- **Ollama** (local): Uses `http://localhost:11434/v1`, no API key required
- **Custom**: Any OpenAI-compatible endpoint (vLLM, LM Studio, etc.) — set `llm_api_base` explicitly
## Available Skills
### enrich_hf_metadata
Enriches HuggingFace-downloaded models with metadata extracted by an LLM from the HF model card.
**Entry point**: Right-click context menu → "Enrich Metadata (Agent)"
**What it does**:
1. Reads the model's `.metadata.json` to get the `hf_url`
2. Fetches the README.md from the HuggingFace repository
3. Sends the README + local metadata to the LLM for structured extraction
4. Writes extracted fields to `.metadata.json`:
- `base_model` — only if current value is empty
- `trainedWords` — trigger words (LoRA only, if none exist)
- `modelDescription` — concise summary (if none exists)
- `tags` — merged with existing tags, deduplicated
- `metadata_source` — audit trail: `agent:enrich_hf_metadata`
- `llm_enriched_at` — ISO timestamp
5. Downloads and optimizes preview image (if LLM found one in the README)
6. Updates the scanner cache
7. Broadcasts WebSocket progress events
**Model types**: LoRA, Checkpoint, Embedding
## Adding a New Skill
### 1. Create the skill directory
```
py/services/agent/skills/<skill_name>/
├── skill.yaml # Skill metadata and schemas
├── prompt.md # LLM prompt template
└── handler.py # Pre-processing and post-processing
```
### 2. Write skill.yaml
```yaml
name: my_skill
title: "My Skill"
description: "What this skill does"
llm_required: true
model_type_filter: ["lora"] # or null for all types
input_schema:
type: object
properties:
model_paths:
type: array
items:
type: string
required:
- model_paths
output_schema:
type: object
properties:
# ... JSON schema for LLM output
permissions:
write_metadata: true
write_previews: false
network_domains:
- "example.com"
```
### 3. Write prompt.md
Use `{{variable}}` placeholders that will be replaced with data from the `prepare` function:
```markdown
You are an expert assistant...
Model URL: {{hf_url}}
README content:
{{readme_content}}
Current metadata:
{{current_metadata}}
```
### 4. Write handler.py
```python
async def prepare(model_path: str, input_data: dict) -> dict:
"""Gather context for the LLM prompt. Returns variables for template rendering."""
return {
"model_path": model_path,
# ... other variables used in prompt.md
}
async def post_process(context) -> dict:
"""Apply the LLM-extracted data to the model."""
llm_response = context.llm_response
# ... write metadata, download previews, update cache
return {
"success": True,
"updated_fields": ["base_model", "tags"],
"errors": [],
}
```
**Important**: Use absolute imports (`from py.utils.metadata_manager import MetadataManager`) because skills are loaded via `importlib.util.spec_from_file_location`, which doesn't support relative imports.
### 5. Test
The skill is automatically discovered by `SkillRegistry` on startup. Test with:
```python
pytest tests/services/test_agent_service.py
```
## API Endpoints
| Method | Path | Description |
|---|---|---|
| GET | `/api/lm/agent/skills` | List available skills |
| POST | `/api/lm/agent/execute/{skill_name}` | Execute a skill (body: `{"model_paths": [...]}`) |
| POST | `/api/lm/agent/cancel` | Cancel running skill (stub) |
## WebSocket Events
| Type | When | Key fields |
|---|---|---|
| `agent_progress` | Skill started/processing | `skill`, `status`, `total`, `processed`, `success`, `current_path` |
| `agent_progress` | Skill completed | `skill`, `status`, `updated_models`, `errors`, `summary` |
| `agent_progress` | Skill error | `skill`, `status`, `error` |
## Security Model
Skills declare permissions in `skill.yaml`:
- `write_metadata` — can write `.metadata.json` files
- `write_previews` — can download/replace preview images
- `network_domains` — allowed domains for HTTP requests
These are declarative constraints checked by `AgentService`. They are defense-in-depth, not a sandbox — the Python process can technically do anything, but the contract is clear and auditable.
## File Locations
| Component | Path |
|---|---|
| LLMService | `py/services/llm_service.py` |
| AgentService | `py/services/agent/agent_service.py` |
| SkillRegistry | `py/services/agent/skill_registry.py` |
| SkillDefinition | `py/services/agent/skill_definition.py` |
| Skills directory | `py/services/agent/skills/` |
| Route handlers | `py/routes/handlers/agent_handlers.py` |
| Frontend manager | `static/js/managers/AgentManager.js` |
| Settings UI | `templates/components/modals/settings_modal.html` |
| Context menu | `templates/components/context_menu.html` |
+65
View File
@@ -0,0 +1,65 @@
# ComfyUI Dual-Mode Widget Rendering
ComfyUI custom node widgets render in one of two modes. Patterns that work in one often fail silently in the other. Test both.
## Mode Detection
```js
typeof LiteGraph !== 'undefined' && LiteGraph.vueNodesMode
```
In Vue SFCs, `window.LiteGraph` is unavailable — pass as a prop from `main.ts`.
## Canvas Mode Layout
Uses `computeLayoutSize()` + `distributeSpace()` to allocate widget height within the node. Widgets with `computeLayoutSize` participate in space distribution; those with `computeSize` have fixed height.
- `getMinHeight()` in `addDOMWidget` options → minimum widget height
- `widget.computeLayoutSize()``{ minHeight, minWidth, maxHeight? }`
- Avoid `getMaxHeight()` unless the widget genuinely needs a fixed cap (prevents user resize)
## Vue Mode Layout
Uses CSS Grid (`grid-template-rows`) + `ResizeObserver`. The ResizeObserver watches the widget's DOM and feeds back into grid row sizing. This creates a feedback loop: content grows → row resizes → more space for content → content reflows/grows → row resizes again.
### Height Containment
The fix: `contain: layout size` on the widget root. This tells the browser the element's intrinsic size is CSS-determined, not driven by descendant content. The ResizeObserver sees a stable size and the loop is broken.
```css
.widget-root.lm-vue-node {
height: 100%;
min-height: var(--comfy-widget-min-height, 200px);
contain: layout size;
}
```
Existing examples: `.lm-loras-container.lm-vue-node` and `.comfy-tags-container.lm-vue-node` in `web/comfyui/lm_styles.css`.
**Do NOT** fix height issues with `maxHeight`, `getMaxHeight()`, or inline `max-height` — these prevent the user from resizing the node.
## Scroll Wheel Isolation
Both modes need to distinguish "user wants to scroll widget content" from "user wants to zoom canvas".
**Canvas mode:** Add `@wheel` on widget root. Check `event.target.closest(selector)` for scrollable sub-areas. If scrollable → `event.stopPropagation()`. Otherwise → `app.canvas.processMouseWheel(event)`.
**Vue mode:** Add CSS class `lm-wheel-scrollable` to scrollable elements. The global capture-phase hook in `web/comfyui/utils.js` (`enableListWheelScroll`) detects wheel events on marked elements and manually scrolls them via `element.scrollTop`, consuming the event before canvas zoom sees it.
## DOM Structure
`main.ts` creates an outer `<div>` container, then `vueApp.mount(container)`. The Vue app renders its own root element inside.
- `container.id` / `container.style.*` → outer element
- Vue scoped `<style>``[data-v-hash]` applies only to Vue root
Classes needed by scoped Vue CSS must go on the Vue root element. Pass data as props and bind with `:class` rather than manipulating the DOM from `main.ts`.
## Serialization
For stateful widgets that need workflow persistence:
- `serialize: true` in `addDOMWidget` options
- `serializeValue()` → state snapshot (called on workflow save)
- `onSetValue(v)` → restore state (called on workflow load)
- Always handle missing keys in restored value for backward compatibility with old workflows
File diff suppressed because one or more lines are too long
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2217 -2019
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+2218 -2020
View File
File diff suppressed because it is too large Load Diff
+68 -8
View File
@@ -8,6 +8,8 @@ from typing import Any, Dict, Iterable, List, Mapping, Optional, Set, Tuple
import logging
import json
import urllib.parse
import sys as _sys
import types as _types
import time
from .utils.cache_paths import CacheType, get_cache_file_path, get_legacy_cache_paths
@@ -175,8 +177,7 @@ class Config:
# Load extra folder paths from active library settings before symlink scan
# so both primary and extra paths are discovered in a single pass.
if not standalone_mode:
self._load_extra_paths_from_settings()
self._load_extra_paths_from_settings()
# Scan symbolic links during initialization
self._initialize_symlink_mappings()
@@ -191,7 +192,7 @@ class Config:
Called during ``Config.__init__`` before the symlink scan so both primary and
extra paths are discovered in a single pass. Mirrors the extra-path
portion of ``_apply_library_paths`` without replacing the primary roots
that were already resolved from ComfyUI's ``folder_paths``.
that were already resolved via ``folder_paths.get_folder_paths``.
"""
try:
from .services.settings_manager import get_settings_manager
@@ -207,6 +208,12 @@ class Config:
if not isinstance(library_config, dict):
return
# Always read recipes_path — it is independent of extra folder paths
# and must be set before any early returns below.
recipes_path = library_config.get("recipes_path", "")
if isinstance(recipes_path, str) and recipes_path:
self.recipes_path = recipes_path
extra_folder_paths = library_config.get("extra_folder_paths")
if not isinstance(extra_folder_paths, dict):
return
@@ -232,10 +239,6 @@ class Config:
extra_embedding
)
recipes_path = library_config.get("recipes_path", "")
if isinstance(recipes_path, str) and recipes_path:
self.recipes_path = recipes_path
if self.extra_loras_roots:
logger.info(
"Found extra LoRA roots:"
@@ -356,6 +359,47 @@ class Config:
"Failed to rename legacy 'default' library: %s", rename_error
)
# Clean up a stale "default" library entry that has no meaningful
# paths configured (e.g. leftover bootstrap artifact). This only
# fires when "comfyui" already exists so we never delete the last
# remaining library.
if (
"default" in libraries
and "comfyui" in libraries
and isinstance(default_library, Mapping)
):
default_folder_paths = _normalize_library_folder_paths(
default_library
)
default_extra_paths = default_library.get("extra_folder_paths", {})
has_meaningful_paths = bool(default_folder_paths) or bool(
default_extra_paths
) or any(
default_library.get(key)
for key in (
"default_lora_root",
"default_checkpoint_root",
"default_unet_root",
"default_embedding_root",
"recipes_path",
)
)
if not has_meaningful_paths:
try:
settings_service.delete_library("default")
libraries_changed = True
logger.info(
"Removed stale 'default' library entry "
"with no meaningful paths configured"
)
libraries = settings_service.get_libraries()
comfy_library = libraries.get("comfyui", {})
except Exception as delete_error:
logger.debug(
"Failed to remove stale 'default' library: %s",
delete_error,
)
default_lora_root = _resolve_valid_default_root(
comfy_library.get("default_lora_root", ""),
list(self.loras_roots or []),
@@ -1380,4 +1424,20 @@ class Config:
# Global config instance
config = Config()
# NOTE: Guard against re-import. When ServiceRegistry.get_lora_scanner() triggers
# a fresh import of lora_scanner → config, we must NOT re-execute Config.__init__()
# (which re-scans all roots, re-registers libraries, etc.).
#
# Strategy: store the config instance in a dedicated sentinel module
# ('_lm_config_cache') that is NEVER removed from sys.modules (its key does
# NOT start with 'py.'), so it survives re-imports of py.* modules.
_CONFIG_SENTINEL = "_lm_config_cache"
if _CONFIG_SENTINEL in _sys.modules:
# Re-import: reuse the existing singleton from the sentinel.
config: Config = _sys.modules[_CONFIG_SENTINEL].config # type: ignore[valid-type]
else:
config: Config = Config()
# Register the sentinel so re-imports of py.config find us.
_sentinel_mod = _types.ModuleType(_CONFIG_SENTINEL)
_sentinel_mod.config = config
_sys.modules[_CONFIG_SENTINEL] = _sentinel_mod
+20
View File
@@ -208,6 +208,10 @@ class LoraManager:
# Initialize WebSocket manager
await ServiceRegistry.get_websocket_manager()
# Preload LLM model catalog (background task, non-blocking)
from .services.llm_service import LLMService
await LLMService.get_instance()
# Initialize scanners in background
lora_scanner = await ServiceRegistry.get_lora_scanner()
checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner()
@@ -436,5 +440,21 @@ class LoraManager:
try:
logger.info("LoRA Manager: Cleaning up services")
# Cancel any in-flight scanner initialization tasks so thread-pool
# workers (e.g. _initialize_cache_sync) can break out of their loops
# when the server shuts down (e.g. Ctrl+C on WSL).
for name in ("lora_scanner", "checkpoint_scanner", "embedding_scanner"):
scanner = ServiceRegistry.get_service_sync(name)
if scanner is not None and hasattr(scanner, "cancel_task"):
scanner.cancel_task()
logger.debug("LoRA Manager: Cancelled %s", name)
# Close shared aiohttp sessions to avoid "Unclosed client session" warnings
try:
from py.routes.handlers.hf_handlers import close_hf_api_session
await close_hf_api_session()
except Exception as exc:
logger.debug("Error closing HF API session: %s", exc)
except Exception as e:
logger.error(f"Error during cleanup: {e}", exc_info=True)
+15 -1
View File
@@ -1,5 +1,11 @@
"""Constants used by the metadata collector"""
# Sentinel value for clip_skip to distinguish "unconnected / widget default"
# from "user wired value 0". Both ComfyUI CLIPSetLastLayer (-24..-1) and
# A1111 conventions treat 0 as meaningless for clip skipping, but users may
# explicitly wire 0 to the overwrite node to express "no clip skip / default".
CLIP_SKIP_SENTINEL = -25
# Metadata categories
MODELS = "models"
PROMPTS = "prompts"
@@ -9,6 +15,14 @@ EMBEDDINGS = "embeddings"
SIZE = "size"
IMAGES = "images"
IS_SAMPLER = "is_sampler" # New constant to mark sampler nodes
OVERWRITE = "overwrite" # Manual metadata overwrite from MetadataOverwriteLM node
# Field names that the MetadataOverwriteLM node and its extractor share
METADATA_OVERWRITE_FIELDS = (
"prompt", "negative_prompt", "seed", "steps", "cfg_scale",
"sampler", "scheduler", "model", "loras", "size",
"clip_skip", "additional_data",
)
# Complete list of categories to track
METADATA_CATEGORIES = [MODELS, PROMPTS, SAMPLING, LORAS, EMBEDDINGS, SIZE, IMAGES]
METADATA_CATEGORIES = [MODELS, PROMPTS, SAMPLING, LORAS, EMBEDDINGS, SIZE, IMAGES, OVERWRITE]
+16 -6
View File
@@ -83,7 +83,8 @@ class MetadataHook:
# Record inputs before execution
if node_id is not None:
registry.record_node_execution(node_id, class_type, input_data_all, None)
return_types = getattr(obj, 'RETURN_TYPES', None)
registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
except Exception as e:
logger.error(f"Error collecting metadata (pre-execution): {str(e)}")
@@ -114,7 +115,8 @@ class MetadataHook:
# Record outputs after execution
if node_id is not None:
registry.update_node_execution(node_id, class_type, results)
return_types = getattr(obj, 'RETURN_TYPES', None)
registry.update_node_execution(node_id, class_type, results, return_types=return_types)
except Exception as e:
logger.error(f"Error collecting metadata (post-execution): {str(e)}")
@@ -135,10 +137,13 @@ class MetadataHook:
# Store the dynprompt reference for node lookups
if hasattr(prompt, 'original_prompt'):
registry.set_current_prompt(prompt)
# Store extra_data for accessing full workflow node properties
registry.set_extra_data(extra_data)
# Execute the original function
return original_execute(*args, **kwargs)
# Replace the functions
execution._map_node_over_list = map_node_over_list_with_metadata
execution.execute = execute_with_prompt_tracking
@@ -163,7 +168,8 @@ class MetadataHook:
class_type = obj.__class__.__name__
node_id = unique_id
if node_id is not None:
registry.record_node_execution(node_id, class_type, input_data_all, None)
return_types = getattr(obj, 'RETURN_TYPES', None)
registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
except Exception as e:
logger.error(f"Error collecting metadata (pre-execution): {str(e)}")
@@ -180,7 +186,8 @@ class MetadataHook:
class_type = obj.__class__.__name__
node_id = unique_id
if node_id is not None:
registry.update_node_execution(node_id, class_type, results)
return_types = getattr(obj, 'RETURN_TYPES', None)
registry.update_node_execution(node_id, class_type, results, return_types=return_types)
except Exception as e:
logger.error(f"Error collecting metadata (post-execution): {str(e)}")
@@ -202,6 +209,9 @@ class MetadataHook:
if hasattr(prompt, 'original_prompt'):
registry.set_current_prompt(prompt)
# Store extra_data for accessing full workflow node properties
registry.set_extra_data(extra_data)
# Execute the original function
return await original_execute(*args, **kwargs)
+138 -14
View File
@@ -1,15 +1,68 @@
import json
import logging
import os
from .constants import IMAGES
# Check if running in standalone mode
standalone_mode = os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1" or os.environ.get("HF_HUB_DISABLE_TELEMETRY", "0") == "0"
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IS_SAMPLER
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IS_SAMPLER, OVERWRITE
from .node_extractors import NODE_EXTRACTORS
logger = logging.getLogger(__name__)
# Keys that identify metadata hint marks stored in node.properties.lm_marker_role
_META_MARK_PREFIX = "meta_"
_MARK_PRIMARY_MODEL = "primary_model"
_MARK_PRIMARY_SAMPLER = "primary_sampler"
_MARK_POSITIVE_PROMPT = "positive_prompt"
_MARK_NEGATIVE_PROMPT = "negative_prompt"
class MetadataProcessor:
"""Process and format collected metadata"""
@staticmethod
def _get_user_marks(metadata):
"""Scan workflow nodes (from extra_data.extra_pnginfo.workflow) for user-assigned
metadata hint marks stored in node.properties.lm_marker_role.
Returns a dict mapping mark type keys to node IDs.
Example: {'primary_model': '42', 'primary_sampler': '17'}
"""
marks: dict[str, str] = {}
# Primary source: extra_data.extra_pnginfo.workflow.nodes (has full properties)
extra_data = metadata.get("extra_data")
if extra_data and isinstance(extra_data, dict):
extra_pnginfo = extra_data.get("extra_pnginfo", {})
if isinstance(extra_pnginfo, dict):
workflow = extra_pnginfo.get("workflow", {})
nodes = workflow.get("nodes", [])
for node in nodes:
node_id = str(node.get("id", ""))
role = node.get("properties", {}).get("lm_marker_role", "")
if role.startswith(_META_MARK_PREFIX):
mark_type = role[len(_META_MARK_PREFIX):]
if mark_type in marks:
logger.warning(
"Duplicate meta hint '%s': node %s (previous: %s), "
"last match wins",
mark_type, node_id, marks[mark_type],
)
marks[mark_type] = node_id
# Fallback: try prompt.original_prompt (API-only submissions may not have workflow)
if not marks:
prompt = metadata.get("current_prompt")
if prompt and getattr(prompt, "original_prompt", None):
for node_id, node_data in prompt.original_prompt.items():
role = node_data.get("properties", {}).get("lm_marker_role", "")
if role.startswith(_META_MARK_PREFIX):
mark_type = role[len(_META_MARK_PREFIX):]
marks[mark_type] = node_id
return marks
@staticmethod
def find_primary_sampler(metadata, downstream_id=None):
"""
@@ -471,20 +524,57 @@ class MetadataProcessor:
"checkpoint": None,
"loras": "",
"size": None,
"clip_skip": None
"clip_skip": None,
"additional_data": "",
}
# Get the prompt object for node relationship tracing
prompt = metadata.get("current_prompt")
# Find the primary KSampler node
primary_sampler_id, primary_sampler = MetadataProcessor.find_primary_sampler(metadata, id)
# Directly get checkpoint from metadata instead of tracing
# Pass primary_sampler_id to avoid redundant calculation
checkpoint = MetadataProcessor.find_primary_checkpoint(metadata, id, primary_sampler_id)
if checkpoint:
params["checkpoint"] = checkpoint
# ---- User marks: override heuristic inference with user-assigned hints ----
user_marks = MetadataProcessor._get_user_marks(metadata)
# Find the primary KSampler node (user mark takes priority)
primary_sampler_id = None
primary_sampler = None
if _MARK_PRIMARY_SAMPLER in user_marks:
marked_id = user_marks[_MARK_PRIMARY_SAMPLER]
sampler_data = metadata.get(SAMPLING, {}).get(marked_id)
if sampler_data and sampler_data.get(IS_SAMPLER):
primary_sampler_id = marked_id
primary_sampler = sampler_data
else:
logger.warning(
"User-marked primary sampler %s has no runtime metadata, "
"falling back to heuristic",
marked_id,
)
if primary_sampler is None:
primary_sampler_id, primary_sampler = MetadataProcessor.find_primary_sampler(metadata, id)
# Resolve checkpoint / model (user mark takes priority)
if _MARK_PRIMARY_MODEL in user_marks:
marked_id = user_marks[_MARK_PRIMARY_MODEL]
if marked_id in metadata.get(MODELS, {}):
params["checkpoint"] = metadata[MODELS][marked_id].get("name")
else:
extra_data = metadata.get("extra_data")
extra_pnginfo = extra_data.get("extra_pnginfo", {}) if extra_data and isinstance(extra_data, dict) else {}
workflow = extra_pnginfo.get("workflow", {}) if isinstance(extra_pnginfo, dict) else {}
node_type = "unknown"
for n in workflow.get("nodes", []):
if str(n.get("id", "")) == marked_id:
node_type = n.get("type", "unknown")
break
logger.warning(
"User-marked primary model %s (type=%s, registered=%s) has no runtime metadata, "
"falling back to heuristic",
marked_id, node_type, node_type in NODE_EXTRACTORS,
)
if params["checkpoint"] is None:
checkpoint = MetadataProcessor.find_primary_checkpoint(metadata, id, primary_sampler_id)
if checkpoint:
params["checkpoint"] = checkpoint
# Check if guidance parameter exists in any sampling node
for node_id, sampler_info in metadata.get(SAMPLING, {}).items():
@@ -539,7 +629,22 @@ class MetadataProcessor:
# For SamplerCustom, handle any additional parameters
MetadataProcessor.handle_custom_advanced_sampler(metadata, prompt, primary_sampler_id, params)
# ---- User marks: override prompts with explicitly tagged nodes ----
prompts_data = metadata.get(PROMPTS, {})
if _MARK_POSITIVE_PROMPT in user_marks:
pos_id = user_marks[_MARK_POSITIVE_PROMPT]
if pos_id in prompts_data:
prompt_text = prompts_data[pos_id].get("text") or prompts_data[pos_id].get("positive_text")
if prompt_text:
params["prompt"] = prompt_text
if _MARK_NEGATIVE_PROMPT in user_marks:
neg_id = user_marks[_MARK_NEGATIVE_PROMPT]
if neg_id in prompts_data:
prompt_text = prompts_data[neg_id].get("text") or prompts_data[neg_id].get("negative_text")
if prompt_text:
params["negative_prompt"] = prompt_text
# Size extraction is same for all sampler types
# Check if the sampler itself has size information (from latent_image)
if primary_sampler_id in metadata.get(SIZE, {}):
@@ -568,7 +673,26 @@ class MetadataProcessor:
break
if params["clip_skip"] is None:
params["clip_skip"] = "1"
# ---- Apply manual metadata overwrites ----
for overwrite_info in metadata.get(OVERWRITE, {}).values():
overwrite_params = overwrite_info.get("parameters", {})
for key, value in overwrite_params.items():
if key == "clip_skip":
# Accept any value from overwrite node (sentinel -25 already
# filtered upstream). Needed because falsy check treats 0
# as "not set" even though 0 is a valid wired input here.
params[key] = value
elif value: # truthy check — only overwrite when user provided a real value
params[key] = value
# Bridge: the overwrite node exposes the field as "model" (more accurate),
# but the internal pipeline key remains "checkpoint" for backward compatibility
# with A1111 metadata format and downstream consumers.
if params.get("model"):
params["checkpoint"] = params["model"]
del params["model"]
return params
@staticmethod
+36 -13
View File
@@ -1,7 +1,7 @@
import time
from nodes import NODE_CLASS_MAPPINGS # type: ignore
from .node_extractors import NODE_EXTRACTORS, GenericNodeExtractor
from .constants import METADATA_CATEGORIES, IMAGES
from .constants import METADATA_CATEGORIES, IMAGES, OVERWRITE
class MetadataRegistry:
@@ -61,6 +61,7 @@ class MetadataRegistry:
{
"execution_order": [],
"current_prompt": None, # Will store the prompt object
"extra_data": None, # Will store the API extra_data for workflow metadata
"timestamp": time.time(),
}
)
@@ -75,6 +76,11 @@ class MetadataRegistry:
# Store the prompt in the metadata for later relationship tracing
self.prompt_metadata[self.current_prompt_id]["current_prompt"] = prompt
def set_extra_data(self, extra_data):
"""Store the API extra_data (contains extra_pnginfo.workflow with node properties)"""
if self.current_prompt_id and self.current_prompt_id in self.prompt_metadata:
self.prompt_metadata[self.current_prompt_id]["extra_data"] = extra_data
def get_metadata(self, prompt_id=None):
"""Get collected metadata for a prompt"""
key = prompt_id if prompt_id is not None else self.current_prompt_id
@@ -122,20 +128,28 @@ class MetadataRegistry:
cache_key = f"{node_id}:{class_type}"
# Check if this node type is relevant for metadata collection
if class_type in NODE_EXTRACTORS:
if class_type in NODE_EXTRACTORS or cache_key in self.node_cache:
# Check if we have cached metadata for this node
if cache_key in self.node_cache:
cached_data = self.node_cache[cache_key]
# Detect bypass (mode=4) / mute (mode=2) — these nodes
# were intentionally disabled and should not contribute
# overwrite values from a previous execution's cache.
node_mode = node_data.get("mode", 0)
node_is_disabled = node_mode in (2, 4)
# Apply cached metadata to the current metadata
for category in self.metadata_categories:
if category == OVERWRITE and node_is_disabled:
continue
if category in cached_data and node_id in cached_data[category]:
if node_id not in metadata[category]:
metadata[category][node_id] = cached_data[category][
node_id
]
def record_node_execution(self, node_id, class_type, inputs, outputs):
def record_node_execution(self, node_id, class_type, inputs, outputs, return_types=None):
"""Record information about a node's execution"""
if not self.current_prompt_id:
return
@@ -158,17 +172,18 @@ class MetadataRegistry:
# Extract node-specific metadata
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
extractor.extract(
node_id,
processed_inputs,
outputs,
self.prompt_metadata[self.current_prompt_id],
)
if extractor is GenericNodeExtractor:
extractor.extract(node_id, processed_inputs, outputs,
self.prompt_metadata[self.current_prompt_id],
return_types=return_types)
else:
extractor.extract(node_id, processed_inputs, outputs,
self.prompt_metadata[self.current_prompt_id])
# Cache this node's metadata
self._cache_node_metadata(node_id, class_type)
def update_node_execution(self, node_id, class_type, outputs):
def update_node_execution(self, node_id, class_type, outputs, return_types=None):
"""Update node metadata with output information"""
if not self.current_prompt_id:
return
@@ -179,9 +194,17 @@ class MetadataRegistry:
# Use the same extractor to update with outputs
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
if hasattr(extractor, "update"):
extractor.update(
node_id, processed_outputs, self.prompt_metadata[self.current_prompt_id]
)
if extractor is GenericNodeExtractor:
extractor.update(
node_id, processed_outputs,
self.prompt_metadata[self.current_prompt_id],
return_types=return_types,
)
else:
extractor.update(
node_id, processed_outputs,
self.prompt_metadata[self.current_prompt_id],
)
# Update the cached metadata for this node
self._cache_node_metadata(node_id, class_type)
+153 -5
View File
@@ -2,7 +2,7 @@ import json
import os
import re
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER
from .constants import CLIP_SKIP_SENTINEL, MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER, OVERWRITE, METADATA_OVERWRITE_FIELDS
def _store_checkpoint_metadata(metadata, node_id, model_name):
@@ -31,11 +31,78 @@ class NodeMetadataExtractor:
pass
class GenericNodeExtractor(NodeMetadataExtractor):
"""Default extractor for nodes without specific handling"""
"""Fallback extractor with type-signature-based detection.
When a node is not in the NODE_EXTRACTORS registry, the hook layer
passes ``return_types`` from ``obj.RETURN_TYPES``:
* ``MODEL`` output: common input fields (ckpt_name, unet_name, etc.)
are checked for a model file name and stored as checkpoint metadata.
* ``CONDITIONING`` output: common text input fields are checked for
prompt text and stored as prompt metadata.
"""
# Input field names that carry a model path in loader-style nodes.
_MODEL_NAME_FIELDS = (
"ckpt_name", "unet_name", "model_path", "model_name", "gguf_name",
)
# Extensions used by checkpoint_scanner.py — only record values that look
# like real model filenames to avoid capturing unrelated string fields.
_MODEL_EXTENSIONS = {
".ckpt", ".pt", ".pt2", ".bin", ".pth", ".safetensors", ".pkl", ".sft", ".gguf",
}
# Input field names that may carry prompt text in encoder-style nodes.
_TEXT_FIELDS = ("text", "clip_l", "t5xxl", "prompt", "positive", "negative")
@staticmethod
def extract(node_id, inputs, outputs, metadata):
pass
def extract(node_id, inputs, outputs, metadata, return_types=None):
if return_types is None:
return
# — MODEL loader detection (checkpoint / UNET / GGUF) —
if "MODEL" in return_types or any("MODEL" in str(t) for t in return_types):
for field in GenericNodeExtractor._MODEL_NAME_FIELDS:
val = inputs.get(field)
if val and isinstance(val, str) and val.strip():
name = val.strip()
if not any(name.lower().endswith(ext) for ext in GenericNodeExtractor._MODEL_EXTENSIONS):
continue
_store_checkpoint_metadata(metadata, node_id, name)
return
# — CONDITIONING encoder detection (CLIPTextEncode, Flux, custom) —
if "CONDITIONING" in return_types or any("CONDITIONING" in str(t) for t in return_types):
text = None
for field in GenericNodeExtractor._TEXT_FIELDS:
val = inputs.get(field)
if val and isinstance(val, str) and val.strip():
text = val.strip()
break
if text:
prompt_data = metadata.setdefault(PROMPTS, {})
prompt_data[node_id] = {
"text": text,
"node_id": node_id,
}
@staticmethod
def update(node_id, outputs, metadata, return_types=None):
if return_types is None:
return
if "CONDITIONING" not in return_types and not any(
"CONDITIONING" in str(t) for t in return_types
):
return
if node_id not in metadata.get(PROMPTS, {}):
return
if outputs and isinstance(outputs, list) and len(outputs) > 0:
if isinstance(outputs[0], tuple) and len(outputs[0]) > 0:
cond = outputs[0][0]
if cond is not None:
metadata[PROMPTS][node_id]["conditioning"] = cond
class CheckpointLoaderExtractor(NodeMetadataExtractor):
@staticmethod
def extract(node_id, inputs, outputs, metadata):
@@ -901,6 +968,55 @@ class LoraLoaderManagerExtractor(NodeMetadataExtractor):
"node_id": node_id
}
class LoraTextLoaderManagerExtractor(NodeMetadataExtractor):
"""Extract LoRA metadata from LoraTextLoaderLM (LoRA Text Loader).
The node accepts a `lora_syntax` STRING containing <lora:name:strength> tags
(same format as the ComfyUI prompt), plus an optional `lora_stack`.
This extractor parses the syntax string using the same regex as the node.
"""
@staticmethod
def extract(node_id, inputs, outputs, metadata):
if not inputs:
return
active_loras = []
# Process lora_stack if available (optional input)
if "lora_stack" in inputs:
lora_stack = inputs.get("lora_stack", [])
for item in lora_stack:
# lora_stack entries are (path, model_strength, clip_strength) tuples
if isinstance(item, (list, tuple)) and len(item) >= 2:
lora_path = item[0]
model_strength = item[1]
lora_name = os.path.splitext(os.path.basename(lora_path))[0]
active_loras.append({
"name": lora_name,
"strength": round(float(model_strength), 2)
})
# Process lora_syntax string input
if "lora_syntax" in inputs:
lora_syntax = inputs.get("lora_syntax", "")
if lora_syntax and isinstance(lora_syntax, str):
pattern = r"<lora:([^:>]+):([^:>]+)(?::([^:>]+))?>"
matches = re.findall(pattern, lora_syntax, re.IGNORECASE)
for match in matches:
lora_name = match[0]
model_strength = float(match[1])
active_loras.append({
"name": lora_name,
"strength": round(model_strength, 2)
})
if active_loras:
metadata[LORAS][node_id] = {
"lora_list": active_loras,
"node_id": node_id
}
class FluxGuidanceExtractor(NodeMetadataExtractor):
@staticmethod
def extract(node_id, inputs, outputs, metadata):
@@ -1105,6 +1221,35 @@ class CR_ApplyControlNetStackExtractor(NodeMetadataExtractor):
metadata[PROMPTS][node_id]["positive_encoded"] = transformed_positive
metadata[PROMPTS][node_id]["negative_encoded"] = transformed_negative
class MetadataOverwriteExtractor(NodeMetadataExtractor):
"""Extract manually specified metadata from MetadataOverwriteLM node.
Stores truthy input values under the OVERWRITE category so that
extract_generation_params can merge them over the inferred params.
"""
@staticmethod
def extract(node_id, inputs, outputs, metadata):
if not inputs:
return
overwrite_params = {}
for key in METADATA_OVERWRITE_FIELDS:
value = inputs.get(key)
if key == "clip_skip":
if value != CLIP_SKIP_SENTINEL:
overwrite_params[key] = value
elif value: # truthy — only overwrite when user provided a real value
overwrite_params[key] = value
if overwrite_params:
metadata.setdefault(OVERWRITE, {})
metadata[OVERWRITE][node_id] = {
"parameters": overwrite_params,
"node_id": node_id,
}
# Registry of node-specific extractors
# Keys are node class names
NODE_EXTRACTORS = {
@@ -1146,6 +1291,7 @@ NODE_EXTRACTORS = {
"UNETLoaderLM": UNETLoaderExtractor, # LoRA Manager
"LoraLoader": LoraLoaderExtractor,
"LoraLoaderLM": LoraLoaderManagerExtractor,
"LoraTextLoaderLM": LoraTextLoaderManagerExtractor,
"RgthreePowerLoraLoader": RgthreePowerLoraLoaderExtractor,
"TensorRTLoader": TensorRTLoaderExtractor,
# Conditioning
@@ -1171,5 +1317,7 @@ NODE_EXTRACTORS = {
"CFGGuider": CFGGuiderExtractor, # Add CFGGuider
# Image
"VAEDecode": VAEDecodeExtractor, # Added VAEDecode extractor
# Metadata overwrite
"MetadataOverwriteLM": MetadataOverwriteExtractor,
# Add other nodes as needed
}
+233
View File
@@ -0,0 +1,233 @@
"""Metadata operations — thin in-process wrappers around LoRA Manager internal services.
All functions are simple Python async functions that delegate to the
appropriate internal service. They use **relative imports** within the
``py`` package, so ``sys.modules`` caching works normally and there is no
risk of double import or circular dependencies.
Usage (in-process, primary)::
from py.metadata_ops import list_base_models, read_metadata
models = await list_base_models()
meta = await read_metadata("/path/to/model.safetensors")
Usage (subprocess, debugging / external)::
python -m py.metadata_ops base-models list
python -m py.metadata_ops metadata read /path/to/model.safetensors
"""
from __future__ import annotations
import asyncio
import logging
import os
from typing import Any, Dict, List, Optional
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
SCANNER_TYPE_MAP: dict[str, str] = {
"get_lora_scanner": "lora",
"get_checkpoint_scanner": "checkpoint",
"get_embedding_scanner": "embedding",
}
SCANNER_GETTER_NAMES = tuple(SCANNER_TYPE_MAP.keys())
async def _find_model_entry(
model_path: str,
) -> tuple[object, object, str | None] | tuple[None, None, None]:
"""Iterate all scanners and return the first (scanner, entry, getter_name)
that owns *model_path*. Returns ``(None, None, None)`` when no scanner
claims it.
"""
from ..services.service_registry import ServiceRegistry
normalized = os.path.normpath(model_path)
for getter_name in SCANNER_GETTER_NAMES:
getter = getattr(ServiceRegistry, getter_name, None)
if getter is None:
continue
try:
scanner = await getter()
if scanner is None:
continue
cache = await scanner.get_cached_data()
for entry in cache.raw_data:
if os.path.normpath(entry.get("file_path", "")) == normalized:
return scanner, entry, getter_name
except Exception as exc:
logger.debug(
"Scanner %s check failed for %s: %s",
getter_name, model_path, exc,
)
return None, None, None
async def _find_scanner_for_model(
model_path: str,
) -> tuple[object, object] | tuple[None, None]:
"""Find the (scanner, cache_entry) responsible for *model_path*."""
scanner, entry, _ = await _find_model_entry(model_path)
return scanner, entry
async def identify_model_type(model_path: str) -> str:
"""Determine the model type (``\"lora\"``, ``\"checkpoint\"``, or
``\"embedding\"``) for *model_path*.
Falls back to ``\"lora\"`` when unknown.
"""
_, _, getter_name = await _find_model_entry(model_path)
return SCANNER_TYPE_MAP[getter_name] if getter_name else "lora"
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
async def list_base_models(limit: int = 0) -> List[str]:
"""Return all valid CivitAI base model names.
Uses ``CivitaiBaseModelService.get_base_models()`` which merges a
hardcoded list (``SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS``) with remote
models fetched from the CivitAI API. Never empty the hardcoded
fallback always provides a complete set.
The result is sorted alphabetically. Pass *limit* = 0 for all models.
"""
from ..services.civitai_base_model_service import (
CivitaiBaseModelService,
)
try:
service = await CivitaiBaseModelService.get_instance()
response = await service.get_base_models()
names: List[str] = response.get("models", [])
except Exception as exc:
logger.warning("list_base_models failed: %s", exc)
names = []
if limit > 0:
return names[:limit]
return names
async def read_metadata(model_path: str) -> Dict[str, Any]:
"""Load the full metadata payload for *model_path* from disk.
Returns an empty dict when the metadata file does not exist or cannot
be parsed never raises.
"""
from ..utils.metadata_manager import MetadataManager
try:
return await MetadataManager.load_metadata_payload(model_path) or {}
except Exception as exc:
logger.warning("read_metadata failed for %s: %s", model_path, exc)
return {}
async def apply_metadata_updates(
model_path: str,
updates: Dict[str, Any],
) -> List[str]:
"""Merge *updates* into the model's on-disk metadata and persist.
Returns the list of field names that actually changed.
"""
from ..utils.metadata_manager import MetadataManager
metadata = await read_metadata(model_path)
updated_fields: List[str] = []
for key, value in updates.items():
old = metadata.get(key)
if old != value:
metadata[key] = value
updated_fields.append(key)
if updated_fields:
await MetadataManager.save_metadata(model_path, metadata)
return updated_fields
async def download_preview(
model_path: str,
url: str,
*,
target_width: int = 480,
quality: int = 85,
) -> str | None:
"""Download a preview image from *url*, optimise to .webp, and save it.
The output file is placed alongside the model file with a ``.webp``
extension. Returns the local file path on success, ``None`` on failure.
"""
from ..services.downloader import get_downloader
from ..utils.exif_utils import ExifUtils
if not url or not url.strip():
return None
base_name = os.path.splitext(os.path.basename(model_path))[0]
preview_dir = os.path.dirname(model_path)
output_path = os.path.join(preview_dir, base_name + ".webp")
downloader = await get_downloader()
# Try in-memory download + optimise first
success, content, _headers = await downloader.download_to_memory(
url, use_auth=False,
)
if success and content:
try:
optimized_data, _ = ExifUtils.optimize_image(
image_data=content,
target_width=target_width,
format="webp",
quality=quality,
preserve_metadata=False,
)
with open(output_path, "wb") as f:
f.write(optimized_data)
return output_path
except Exception as exc:
logger.warning("Preview optimisation failed, saving raw: %s", exc)
# Fall through to raw save
# Fallback: download directly to file
try:
ok, _ = await downloader.download_file(url, output_path, use_auth=False)
if ok:
return output_path
except Exception as exc:
logger.warning("Preview fallback download failed for %s: %s", model_path, exc)
return None
async def refresh_cache(model_path: str) -> bool:
"""Invalidate and reload the scanner cache entry for *model_path*.
Returns ``True`` when the model was found and the cache was refreshed.
"""
scanner, entry = await _find_scanner_for_model(model_path)
if scanner is None:
logger.warning("refresh_cache: no scanner found for %s", model_path)
return False
try:
metadata = await read_metadata(model_path)
if not metadata:
logger.warning("refresh_cache: no metadata for %s", model_path)
return False
await scanner.update_single_model_cache(model_path, model_path, metadata)
return True
except Exception as exc:
logger.warning("refresh_cache failed for %s: %s", model_path, exc)
return False
+113
View File
@@ -0,0 +1,113 @@
"""Subprocess entry point for ``metadata_ops`` (debugging / external use).
Usage::
python -m py.metadata_ops base-models list [--limit N]
python -m py.metadata_ops metadata read <path>
python -m py.metadata_ops metadata update <path> --json '{...}'
python -m py.metadata_ops preview download <path> --url <url>
python -m py.metadata_ops cache refresh <path>
"""
from __future__ import annotations
import argparse
import asyncio
import json
import sys
from typing import Any, Dict, List
def _build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(prog="lmcli", description="LoRA Manager Agent CLI")
sub = parser.add_subparsers(dest="command", required=True)
# base-models list
base_models = sub.add_parser("base-models", aliases=["bm"])
base_models_cmds = base_models.add_subparsers(dest="subcommand", required=True)
base_models_list = base_models_cmds.add_parser("list")
base_models_list.add_argument(
"--limit", type=int, default=0, help="Max number of models (0 = all)"
)
# metadata read
meta = sub.add_parser("metadata", aliases=["md"])
meta_cmds = meta.add_subparsers(dest="subcommand", required=True)
meta_read = meta_cmds.add_parser("read")
meta_read.add_argument("path", type=str, help="Model file path")
# metadata update
meta_update = meta_cmds.add_parser("update")
meta_update.add_argument("path", type=str, help="Model file path")
meta_update.add_argument(
"--json",
type=str,
required=True,
help='JSON object of fields to update, e.g. \'{"base_model": "SDXL 1.0"}\'',
)
# preview download
prev = sub.add_parser("preview", aliases=["pv"])
prev_cmds = prev.add_subparsers(dest="subcommand", required=True)
prev_dl = prev_cmds.add_parser("download")
prev_dl.add_argument("path", type=str, help="Model file path")
prev_dl.add_argument("--url", type=str, required=True, help="Preview image URL")
# cache refresh
cache = sub.add_parser("cache")
cache_cmds = cache.add_subparsers(dest="subcommand", required=True)
cache_refresh = cache_cmds.add_parser("refresh")
cache_refresh.add_argument("path", type=str, help="Model file path")
return parser
async def _run(args: argparse.Namespace) -> Any:
from . import ( # lazy import so startup is fast
list_base_models,
read_metadata,
apply_metadata_updates,
download_preview,
refresh_cache,
)
cmd = args.command
sub = args.subcommand
if cmd in ("base-models", "bm") and sub == "list":
return await list_base_models(limit=args.limit)
if cmd in ("metadata", "md") and sub == "read":
return await read_metadata(args.path)
if cmd in ("metadata", "md") and sub == "update":
updates: Dict[str, Any] = json.loads(args.json)
return await apply_metadata_updates(args.path, updates)
if cmd in ("preview", "pv") and sub == "download":
return await download_preview(args.path, args.url)
if cmd == "cache" and sub == "refresh":
return await refresh_cache(args.path)
raise ValueError(f"Unknown command: {cmd} {sub}")
def main() -> None:
parser = _build_parser()
args = parser.parse_args()
result = asyncio.run(_run(args))
# Always print as JSON so callers can parse reliably
if isinstance(result, list):
for item in result:
print(item)
elif isinstance(result, dict):
json.dump(result, sys.stdout, ensure_ascii=False, indent=2)
print()
else:
print(json.dumps(result))
if __name__ == "__main__":
main()
+2
View File
@@ -16,6 +16,8 @@ IMG_EXTENSIONS = (
".tif",
".tiff",
".webp",
".avif",
".jxl",
".mp4"
)
+6 -1
View File
@@ -41,7 +41,12 @@ async def api_json_error(
if exc.status < 400:
raise
logger.warning(
# Preview 404 is routine (file deleted from disk) — not worth a warning.
logger_method = logger.warning
if request.path.startswith("/api/lm/previews") and exc.status == 404:
logger_method = logger.debug
logger_method(
"API %s %s returned HTTP %d: %s",
request.method,
request.path,
+117
View File
@@ -0,0 +1,117 @@
"""Create Hook LoRA (LoraManager) — multi-LoRA hook node compatible with ComfyUI's built-in hook pipeline.
Produces ``("HOOKS",)`` output that chains seamlessly with downstream hook consumers
(ConditioningSetProperties, SetHookKeyframes, CombineHooks, SetClipHooks, etc.).
"""
from __future__ import annotations
import logging
import os
from ..utils.utils import get_lora_info_absolute
from .utils import (
FlexibleOptionalInputType,
any_type,
apply_lora_syntax_format,
get_loras_list,
)
logger = logging.getLogger(__name__)
class CreateHookLoraLM:
NAME = "Create Hook LoRA (LoraManager)"
CATEGORY = "Lora Manager/hooks"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": (
"AUTOCOMPLETE_TEXT_LORAS",
{
"placeholder": "Search LoRAs to add...",
"tooltip": (
"Search and select LoRAs. Each LoRA gets its own "
"model/clip strength. Hooks chain with prev_hooks."
),
},
),
},
"optional": FlexibleOptionalInputType(any_type),
}
RETURN_TYPES = ("HOOKS", "STRING", "STRING")
RETURN_NAMES = ("HOOKS", "trigger_words", "active_loras")
FUNCTION = "create_hook"
def create_hook(self, text: str, **kwargs):
"""Create a HookGroup from the selected LoRAs, chained with prev_hooks.
Each active LoRA from the widget is loaded and wrapped in a WeightHook
via :func:`comfy.hooks.create_hook_lora`. All hooks are combined into a
single group and returned alongside trigger words and a human-readable
summary of the active LoRAs.
"""
del text # used by the frontend widget only
# Lazy imports: comfy is not available in CI/test environment at module level
import comfy.hooks # type: ignore # noqa: C0415
import comfy.utils # type: ignore # noqa: C0415
prev_hooks: comfy.hooks.HookGroup | None = kwargs.get("prev_hooks")
hook_group = prev_hooks.clone() if prev_hooks is not None else comfy.hooks.HookGroup()
all_trigger_words: list[str] = []
active_loras: list[tuple[str, float, float]] = []
for lora in get_loras_list(kwargs):
if not lora.get("active", False):
continue
lora_name = apply_lora_syntax_format(lora["name"])
model_strength = float(lora["strength"])
clip_strength = float(lora.get("clipStrength", model_strength))
# Skip useless no-op entries (both strengths are zero)
if model_strength == 0.0 and clip_strength == 0.0:
continue
lora_path, trigger_words = get_lora_info_absolute(lora_name)
if not lora_path or not os.path.isfile(lora_path):
logger.warning("LoRA '%s' not found — skipping", lora_name)
continue
try:
lora_weights = comfy.utils.load_torch_file(lora_path, safe_load=True)
lora_hooks = comfy.hooks.create_hook_lora(
lora=lora_weights,
strength_model=model_strength,
strength_clip=clip_strength,
)
except Exception:
logger.exception("Failed to load LoRA '%s' — skipping", lora_name)
continue
hook_group = hook_group.clone_and_combine(lora_hooks)
active_loras.append((lora_name, model_strength, clip_strength))
all_trigger_words.extend(trigger_words)
# Format trigger words (group mode separator)
trigger_words_text = ",, ".join(all_trigger_words) if all_trigger_words else ""
# Format active LoRAs summary
formatted_loras = []
for name, model_s, clip_s in active_loras:
if abs(model_s - clip_s) > 0.001:
formatted_loras.append(
f"<lora:{name}:{model_s}:{clip_s}>"
)
else:
formatted_loras.append(f"<lora:{name}:{model_s}>")
active_loras_text = " ".join(formatted_loras)
return (hook_group, trigger_words_text, active_loras_text)
+45
View File
@@ -0,0 +1,45 @@
"""Lora Info display node — pure frontend node for showing selected LoRA info.
This node does NOT participate in workflow execution. Its single optional
"lora_source" input exists solely as a wire-connection anchor so that the
frontend can traverse the graph and push selection data to connected info nodes.
"""
from __future__ import annotations
class LoraInfoLM:
"""Display node that shows filename and notes for the selected LoRA."""
NAME = "Lora Info (LoraManager)"
CATEGORY = "Lora Manager/utils"
DESCRIPTION = (
"Displays information (filename, notes) about the currently selected "
"LoRA. Connect any output from a LoRA Loader or Stacker to the "
"lora_source input, then select a LoRA in the source widget — the "
"info updates automatically. Does not affect workflow execution."
)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {},
}
RETURN_TYPES = ()
RETURN_NAMES = ()
OUTPUT_NODE = False
FUNCTION = "noop"
def noop(self, **kwargs):
# This node is display-only — no workflow execution needed.
return ()
NODE_CLASS_MAPPINGS = {
LoraInfoLM.NAME: LoraInfoLM,
}
NODE_DISPLAY_NAME_MAPPINGS = {
LoraInfoLM.NAME: "Lora Info (LoraManager)",
}
+2 -17
View File
@@ -1,6 +1,5 @@
import importlib
import logging
import re
import comfy.sd # type: ignore
import comfy.utils # type: ignore
@@ -14,6 +13,7 @@ from .utils import (
extract_lora_name,
get_loras_list,
nunchaku_load_lora,
parse_lora_syntax,
)
logger = logging.getLogger(__name__)
@@ -189,25 +189,10 @@ class LoraTextLoaderLM:
RETURN_NAMES = ("MODEL", "CLIP", "trigger_words", "loaded_loras")
FUNCTION = "load_loras_from_text"
def parse_lora_syntax(self, text):
"""Parse LoRA syntax from text input."""
pattern = r"<lora:([^:>]+):([^:>]+)(?::([^:>]+))?>"
matches = re.findall(pattern, text, re.IGNORECASE)
loras = []
for match in matches:
model_strength = float(match[1])
loras.append({
"name": match[0],
"model_strength": model_strength,
"clip_strength": float(match[2]) if match[2] else model_strength,
})
return loras
def load_loras_from_text(self, model, lora_syntax, clip=None, lora_stack=None):
"""Load LoRAs based on text syntax input."""
lora_entries = _collect_stack_entries(lora_stack)
for lora in self.parse_lora_syntax(lora_syntax):
for lora in parse_lora_syntax(lora_syntax):
lora_path, trigger_words = get_lora_info_absolute(lora["name"])
lora_entries.append({
"name": lora["name"],
+86 -10
View File
@@ -1,26 +1,102 @@
from __future__ import annotations
import inspect
import re
from typing import Any
_STACK_INPUT_PATTERN = re.compile(r"^lora_stack(?:_([ab])|(\d+))$")
def _is_stack_input(name: str) -> bool:
return bool(_STACK_INPUT_PATTERN.match(name))
def _stack_slot_number(name: str) -> int:
"""Numeric slot used to order stack inputs; legacy a/b map to 1/2."""
match = _STACK_INPUT_PATTERN.match(name)
if not match:
return -1
letter, digits = match.group(1), match.group(2)
if digits is not None:
return int(digits)
return 1 if letter == "a" else 2
class _LoraStackOptionalInputs:
"""Lookup that preserves explicit optional inputs and dynamic lora_stack slots."""
def __init__(self, explicit_inputs: dict[str, tuple[str, dict[str, Any]]]) -> None:
self._explicit_inputs = explicit_inputs
def __contains__(self, item: object) -> bool:
if not isinstance(item, str):
return False
return item in self._explicit_inputs or _is_stack_input(item)
def __getitem__(self, key: str) -> tuple[str, dict[str, Any]]:
if key in self._explicit_inputs:
return self._explicit_inputs[key]
if _is_stack_input(key):
return (
"LORA_STACK",
{
"tooltip": "A LoRA stack to combine. Connect to add more inputs.",
},
)
raise KeyError(key)
class LoraStackCombinerLM:
NAME = "Lora Stack Combiner (LoraManager)"
CATEGORY = "Lora Manager/stackers"
DESCRIPTION = (
"Combines multiple LoRA stacks into a single stack. "
"Supports dynamic inputs: connect a stack to add more inputs."
)
@classmethod
def INPUT_TYPES(cls):
optional_inputs: dict[str, tuple[str, dict[str, Any]]] = {
"lora_stack1": (
"LORA_STACK",
{
"tooltip": "A LoRA stack to combine. Connect to add more inputs.",
},
),
"lora_stack2": (
"LORA_STACK",
{
"tooltip": "A LoRA stack to combine. Connect to add more inputs.",
},
),
}
stack = inspect.stack()
if len(stack) > 2 and stack[2].function == "get_input_info":
optional_inputs = _LoraStackOptionalInputs(optional_inputs) # type: ignore[assignment]
return {
"required": {
"lora_stack_a": ("LORA_STACK",),
"lora_stack_b": ("LORA_STACK",),
},
"required": {},
"optional": optional_inputs,
}
RETURN_TYPES = ("LORA_STACK",)
RETURN_NAMES = ("LORA_STACK",)
FUNCTION = "combine_stacks"
def combine_stacks(self, lora_stack_a, lora_stack_b):
combined_stack = []
def combine_stacks(self, lora_stack1=None, lora_stack2=None, **kwargs):
stacks = {
"lora_stack1": lora_stack1,
"lora_stack2": lora_stack2,
}
for key, value in kwargs.items():
if _is_stack_input(key) and value is not None:
stacks[key] = value
if lora_stack_a:
combined_stack.extend(lora_stack_a)
if lora_stack_b:
combined_stack.extend(lora_stack_b)
combined_stack = []
for key in sorted(stacks, key=_stack_slot_number):
stack = stacks[key]
if stack:
combined_stack.extend(stack)
return (combined_stack,)
+62
View File
@@ -0,0 +1,62 @@
"""Node to resolve `<lora:name:strength>` syntax to absolute file system paths.
Takes the loaded_loras / active_loras STRING output from LoraLoaderLM or
LoraStackerLM and resolves each lora name to its absolute path on disk via
the scanner cache. Unknown names are returned as-is.
"""
import logging
from ..utils.utils import get_lora_info_absolute
from .utils import parse_lora_syntax
logger = logging.getLogger(__name__)
class LoraSyntaxToPath:
NAME = "LoRA Syntax → Path (LoraManager)"
CATEGORY = "Lora Manager/utils"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"lora_syntax": (
"STRING",
{
"forceInput": True,
"multiline": True,
"tooltip": (
"<lora:name:strength> formatted text from "
"loaded_loras / active_loras output"
),
},
),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("paths",)
FUNCTION = "resolve"
def resolve(self, lora_syntax: str) -> tuple[str]:
"""Parse <lora:...> syntax and resolve each name to its absolute path."""
if not lora_syntax or not lora_syntax.strip():
logger.info("Received empty lora_syntax input")
return ("",)
parsed = parse_lora_syntax(lora_syntax)
if not parsed:
logger.info("No valid <lora:...> entries found in input")
return ("",)
paths: list[str] = []
for entry in parsed:
try:
absolute_path, _ = get_lora_info_absolute(entry["name"])
paths.append(absolute_path)
except Exception:
logger.warning("Failed to resolve lora '%s', skipping", entry["name"])
continue
return ("\n".join(paths),)
+170
View File
@@ -0,0 +1,170 @@
"""Metadata Overwrite node — allows users to manually specify generation parameters
that override the automatically collected/inferred metadata.
Most inputs have falsy defaults (empty string / 0) which are skipped.
clip_skip uses a sentinel default (-25) so that a wired value of 0 is
preserved both ComfyUI and A1111 conventions have no meaningful 0 value,
but users may wire 0 to express "no clip skip / default".
"""
from typing import Any
from ..metadata_collector.constants import (
CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL,
METADATA_OVERWRITE_FIELDS,
)
class MetadataOverwriteLM:
NAME = "Metadata Overwrite (LoraManager)"
CATEGORY = "Lora Manager/utils"
DESCRIPTION = (
"Manually specify generation parameters to override automatically collected "
"metadata. Only filled/connected inputs will take effect — empty defaults "
"are ignored."
)
@classmethod
def INPUT_TYPES(cls) -> dict[str, Any]:
return {
"optional": {
"prompt": (
"STRING",
{
"default": "",
"multiline": True,
"tooltip": "Positive prompt. Only overwrites when non-empty.",
},
),
"negative_prompt": (
"STRING",
{
"default": "",
"multiline": True,
"tooltip": "Negative prompt. Only overwrites when non-empty.",
},
),
"seed": (
"INT",
{
"default": 0,
"min": 0,
"max": 0xFFFFFFFFFFFFFFFF,
"control_after_generate": False,
"tooltip": "Seed value. Only overwrites when > 0.",
},
),
"steps": (
"INT",
{
"default": 0,
"min": 0,
"max": 10000,
"tooltip": "Number of steps. Only overwrites when > 0.",
},
),
"cfg_scale": (
"FLOAT",
{
"default": 0.0,
"min": 0.0,
"max": 100.0,
"tooltip": "CFG scale. Only overwrites when > 0.",
},
),
"sampler": (
"STRING",
{
"default": "",
"tooltip": "Sampler name. Only overwrites when non-empty.",
},
),
"scheduler": (
"STRING",
{
"default": "",
"tooltip": "Scheduler name. Only overwrites when non-empty.",
},
),
"model": (
"STRING",
{
"default": "",
"tooltip": (
"The checkpoint or diffusion model (UNet) used "
"for generation. Only overwrites when non-empty."
),
},
),
"loras": (
"STRING",
{
"default": "",
"multiline": True,
"tooltip": (
"LoRA syntax, e.g. <lora:name:strength> "
"or <lora:name:model_strength:clip_strength>, "
"separated by spaces. Only overwrites when non-empty."
),
},
),
"size": (
"STRING",
{
"default": "",
"tooltip": (
"Image size in WIDTHxHEIGHT format (e.g. 512x768). "
"Only overwrites when non-empty."
),
},
),
"clip_skip": (
"INT",
{
"default": _CLIP_SKIP_SENTINEL,
"min": -25,
"max": 24,
"tooltip": (
"Clip skip (ComfyUI: -24..-1, A1111: 1+). "
"Default -25 means not set — any other value "
"overwrites."
),
},
),
"additional_data": (
"STRING",
{
"default": "",
"multiline": True,
"tooltip": (
"Additional data to embed in the image metadata. "
"Inserted between Clip skip and Model hash in the "
"A1111-compatible parameters string. "
'Example: "Copyright": "Some license info"'
),
},
),
},
}
RETURN_TYPES = ("METADATA",)
RETURN_NAMES = ("metadata",)
FUNCTION = "collect_metadata"
OUTPUT_NODE = True
def collect_metadata(self, **kwargs: Any) -> tuple[dict[str, Any]]:
"""Collect non-default input values into a metadata dict.
For most fields, a falsy value (empty string, 0) means "not set"
and is skipped. clip_skip uses a dedicated sentinel (-25) so that
a wired value of 0 is preserved and reaches the metadata pipeline.
"""
result: dict[str, Any] = {}
for key in METADATA_OVERWRITE_FIELDS:
value = kwargs.get(key)
if key == "clip_skip":
if value != _CLIP_SKIP_SENTINEL:
result[key] = value
elif value:
result[key] = value
return (result,)
+353 -129
View File
@@ -16,6 +16,156 @@ from PIL import Image, PngImagePlugin
import piexif
import logging
# Civitai-compatible sampler name mapping: ComfyUI internal → A1111 display name
CIVITAI_SAMPLER_MAP = {
"euler": "Euler",
"euler_ancestral": "Euler a",
"lms": "LMS",
"heun": "Heun",
"dpm_2": "DPM2",
"dpm_2_ancestral": "DPM2 a",
"dpmpp_2s_ancestral": "DPM++ 2S a",
"dpmpp_2m": "DPM++ 2M",
"dpmpp_sde": "DPM++ SDE",
"dpmpp_sde_gpu": "DPM++ SDE",
"dpmpp_2m_sde": "DPM++ 2M SDE",
"dpmpp_2m_sde_gpu": "DPM++ 2M SDE",
"dpmpp_3m_sde": "DPM++ 3M SDE",
"dpm_fast": "DPM fast",
"dpm_adaptive": "DPM adaptive",
"ddim": "DDIM",
"plms": "PLMS",
"uni_pc_bh2": "UniPC",
"uni_pc": "UniPC",
"lcm": "LCM",
}
# Base model display name → AIR URN slug
# Sourced from civitai source: src/shared/constants/basemodel.constants.ts
BASE_MODEL_AIR_SLUG = {
# Stable Diffusion family
"SD 1.4": "sd1",
"SD 1.5": "sd1",
"SD 1.5 LCM": "sd1",
"SD 1.5 Hyper": "sd1",
"SD 2.0": "sd2",
"SD 2.0 768": "sd2",
"SD 2.1": "sd2",
"SD 2.1 768": "sd2",
"SD 2.1 Unclip": "sd2",
"SD 3.0": "sd3",
"SD 3.5": "sd35",
"SD 3.5 Large": "sd35",
"SD 3.5 Large Turbo": "sd35",
"SD 3.5 Medium": "sd35",
"SDXL 0.9": "sdxl",
"SDXL 1.0": "sdxl",
"SDXL 1.0 LCM": "sdxl",
"SDXL Lightning": "sdxl",
"SDXL Hyper": "sdxl",
"SDXL Turbo": "sdxl",
"SDXL Distilled": "sdxldistilled",
"Stable Cascade": "scascade",
"Stable Video Diffusion": "svd",
"SVD": "svd",
"SVD XT": "svdxt",
# SDXL community fine-tunes
"Pony": "pony",
"Pony Diffusion": "pony",
"Illustrious": "illustrious",
"NoobAI": "noobai",
"Animagine": "illustrious",
# Flux family
"Flux.1": "flux1",
"Flux.1 D": "flux1",
"Flux.1 S": "flux1",
"Flux.1 Krea": "fluxkrea",
"Flux.1 Kontext": "flux1kontext",
"Flux.2": "flux2",
"Flux.2 D": "flux2",
"Flux.2 Klein 9B": "flux2klein_9b",
"Flux.2 Klein 9B Base": "flux2klein_9b_base",
"Flux.2 Klein 4B": "flux2klein_4b",
"Flux.2 Klein 4B Base": "flux2klein_4b_base",
# Other image models (sorted alphabetically)
"AuraFlow": "auraflow",
"Chroma": "chroma",
"HiDream": "hidream",
"HiDream-O1": "hidream-o1",
"Hunyuan DiT": "hydit1",
"Hunyuan Video": "hyv1",
"Kolors": "kolors",
"Lumina": "lumina",
"Mochi": "mochi",
"ODOR": "odor",
"PixArt Alpha": "pixarta",
"PixArt Sigma": "pixarte",
"Playground v2": "playgroundv2",
"Playground v2.5": "playgroundv2",
"Pony Diffusion V7": "ponyv7",
# Video models
"CogVideoX": "cogvideox",
"LTX Video": "ltxv",
"LTX Video 2": "ltxv2",
"LTX Video 2.3": "ltxv23",
"Wan Video": "wanvideo",
"Wan Video 1.3B T2V": "wanvideo_13b_t2v",
"Wan Video 14B T2V": "wanvideo_14b_t2v",
"Wan Video 14B I2V 480p": "wanvideo_14b_i2v_480p",
"Wan Video 14B I2V 720p": "wanvideo_14b_i2v_720p",
# Third-party / proprietary image models
"Boogu": "boogu",
"Ernie": "ernie",
"Grok": "grok",
"HappyHorse": "happyhorse",
"Ideogram": "ideogram",
"Ideogram 4.0": "ideogram",
"Imagen": "imagen4",
"Imagen 4": "imagen4",
"Krea": "krea2",
"Krea 2": "krea2",
"Lens": "lens",
"MAI": "mai",
"Nano Banana": "nanobanana",
"OpenAI": "openai",
"Reve": "reve",
"Reve 2": "reve",
"Reve 2.1": "reve",
"Seedream": "seedream",
"Sora": "sora2",
"Sora 2": "sora2",
"Veo": "veo3",
"Veo 2": "veo3",
"Veo 3": "veo3",
"ZImageTurbo": "zimageturbo",
"ZImageBase": "zimagebase",
"ZImage": "zimagebase",
# Third-party video models
"Hailuo by MiniMax": "minimax",
"Haiper": "haiper",
"Kling": "kling",
"Lightricks": "lightricks",
"Seedance": "seedance",
"Vidu": "vidu",
# Qwen family
"Qwen": "qwen",
"Qwen 2": "qwen2",
# Anima
"Anima": "anima",
# Special
"Upscaler": "upscaler",
"Other": "other",
}
logger = logging.getLogger(__name__)
@@ -70,11 +220,29 @@ class SaveImageLM:
"tooltip": "Compression quality for JPEG and lossy WebP formats (1-100). Higher values mean better quality but larger files.",
},
),
"webp_method": (
"INT",
{
"default": 6,
"min": 0,
"max": 6,
"tooltip": "WebP compression method (0-6). 0=fastest/largest, 6=slowest/smallest. Only applies when file_format is 'webp'.",
},
),
"jpeg_subsampling": (
"INT",
{
"default": 0,
"min": 0,
"max": 2,
"tooltip": "JPEG chroma subsampling level. 0=4:4:4 (best quality), 1=4:2:2, 2=4:2:0 (smallest files). Only applies when file_format is 'jpeg'.",
},
),
"embed_workflow": (
"BOOLEAN",
{
"default": False,
"tooltip": "Embeds the complete workflow data into the image metadata. Only works with PNG and WebP formats.",
"tooltip": "When enabled, saved images store the complete workflow. Drag the image back into ComfyUI to restore the original node graph. PNG and WebP only.",
},
),
"save_with_metadata": (
@@ -142,148 +310,194 @@ class SaveImageLM:
return None
def format_metadata(self, metadata_dict):
"""Format metadata in the requested format similar to userComment example"""
if not metadata_dict:
return ""
def _resolve_model_cache_entry(self, scanner_type: str, name: str):
"""Resolve model hash, civitai metadata, and base_model from scanner cache.
Returns (hash_str, civitai_dict, base_model_str). All values are empty defaults when not found."""
scanner = ServiceRegistry.get_service_sync(scanner_type)
if scanner is None or not name:
return "", {}, ""
# Helper function to only add parameter if value is not None
def add_param_if_not_none(param_list, label, value):
if value is not None:
param_list.append(f"{label}: {value}")
entry = self._get_cached_model_by_name(scanner, name)
if entry is None:
basename = os.path.splitext(os.path.basename(name))[0]
hash_val = scanner.get_hash_by_filename(basename)
return (hash_val or "").lower(), {}, ""
hash_val = (entry.get("sha256") or "").lower()
civitai = entry.get("civitai") or {}
base_model = entry.get("base_model") or ""
return hash_val, civitai, base_model
@staticmethod
def _get_civitai_sampler_name(sampler_name: str, scheduler: str) -> str:
if sampler_name in CIVITAI_SAMPLER_MAP:
civitai_name = CIVITAI_SAMPLER_MAP[sampler_name]
if scheduler == "karras":
civitai_name += " Karras"
elif scheduler == "exponential":
civitai_name += " Exponential"
return civitai_name
else:
if scheduler and scheduler != "normal":
return f"{sampler_name}_{scheduler}"
return sampler_name
@staticmethod
def _build_air_string(base_model: str, model_type: str, model_id: int, version_id: int) -> str:
slug = BASE_MODEL_AIR_SLUG.get(base_model, "other")
type_lower = model_type.lower() if model_type else "other"
return f"urn:air:{slug}:{type_lower}:civitai:{model_id}@{version_id}"
def format_metadata(self, metadata_dict: dict) -> str:
"""Format metadata as A1111-compatible parameters string with Hashes JSON and Civitai resources."""
if not metadata_dict: return ""
# Extract the prompt and negative prompt
prompt = metadata_dict.get("prompt", "")
negative_prompt = metadata_dict.get("negative_prompt", "")
# Extract loras from the prompt if present
steps = metadata_dict.get("steps")
cfg = metadata_dict.get("guidance")
if cfg is None:
cfg = metadata_dict.get("cfg_scale")
if cfg is None:
cfg = metadata_dict.get("cfg")
seed = metadata_dict.get("seed")
size = metadata_dict.get("size")
sampler = metadata_dict.get("sampler") or ""
scheduler = metadata_dict.get("scheduler") or "normal"
checkpoint = metadata_dict.get("checkpoint") or ""
loras_text = metadata_dict.get("loras", "")
lora_hashes = {}
clip_skip = metadata_dict.get("clip_skip")
# If loras are found, add them on a new line after the prompt
# Parse LoRA entries from <lora:name:strength> format
lora_entries: list[tuple[str, float]] = []
if loras_text:
prompt_with_loras = f"{prompt}\n{loras_text}"
for match in re.findall(r"<lora:([^:]+):([^>]+)>", loras_text):
lora_name, strength_str = match
try:
strength = float(strength_str)
except (ValueError, TypeError):
strength = 1.0
lora_entries.append((lora_name, strength))
# Extract lora names from the format <lora:name:strength>
lora_matches = re.findall(r"<lora:([^:]+):([^>]+)>", loras_text)
# Resolve checkpoint hash and Civitai data from local cache
ckpt_hash, ckpt_civitai, ckpt_base_model = "", {}, ""
ckpt_display_name = ""
if checkpoint:
ckpt_hash, ckpt_civitai, ckpt_base_model = self._resolve_model_cache_entry(
"checkpoint_scanner", checkpoint
)
ckpt_display_name = os.path.splitext(os.path.basename(checkpoint))[0]
# Get hash for each lora
for lora_name, strength in lora_matches:
hash_value = self.get_lora_hash(lora_name)
if hash_value:
lora_hashes[lora_name] = hash_value
else:
prompt_with_loras = prompt
# Resolve LoRA hash and Civitai data from local cache
loras_data: list[dict] = []
for lora_name, strength in lora_entries:
lora_hash, lora_civitai, lora_base_model = self._resolve_model_cache_entry(
"lora_scanner", lora_name
)
loras_data.append({
"name": lora_name,
"strength": strength,
"hash": lora_hash,
"civitai": lora_civitai,
"base_model": lora_base_model,
})
# Format the first part (prompt and loras)
metadata_parts = [prompt_with_loras]
# Build Hashes JSON (A1111 / Civitai standard format)
hashes: dict[str, str] = {}
if ckpt_hash:
hashes["model"] = ckpt_hash[:10].upper()
for lora in loras_data:
if lora["hash"]:
hashes[f"LORA:{lora['name']}"] = lora["hash"][:10].upper()
# Add negative prompt
# Build Civitai resources JSON array
civitai_resources: list[dict] = []
if ckpt_civitai.get("id", 0) > 0:
ckpt_resource: dict = {}
ckpt_type = (ckpt_civitai.get("model") or {}).get("type", "Checkpoint")
model_id = ckpt_civitai.get("modelId", 0)
version_id = ckpt_civitai.get("id", 0)
if model_id and version_id:
ckpt_resource["air"] = self._build_air_string(
ckpt_base_model, ckpt_type, int(model_id), int(version_id)
)
elif version_id:
ckpt_resource["modelVersionId"] = int(version_id)
if ckpt_civitai.get("name"):
ckpt_resource["versionName"] = ckpt_civitai["name"]
if ckpt_resource:
civitai_resources.append(ckpt_resource)
for lora in loras_data:
lora_civitai = lora["civitai"]
if not lora_civitai or lora_civitai.get("id", 0) <= 0:
continue
lora_resource: dict = {"weight": lora["strength"]}
lora_type = (lora_civitai.get("model") or {}).get("type", "LORA")
model_id = lora_civitai.get("modelId", 0)
version_id = lora_civitai.get("id", 0)
if model_id and version_id:
lora_resource["air"] = self._build_air_string(
lora["base_model"], lora_type, int(model_id), int(version_id)
)
elif version_id:
lora_resource["modelVersionId"] = int(version_id)
if lora_civitai.get("name"):
lora_resource["versionName"] = lora_civitai["name"]
civitai_resources.append(lora_resource)
sampler_name = CIVITAI_SAMPLER_MAP.get(sampler, sampler) if sampler else None
scheduler_mapping = {
"normal": "Normal",
"karras": "Karras",
"exponential": "Exponential",
"sgm_uniform": "SGM Uniform",
"sgm_quadratic": "SGM Quadratic",
}
scheduler_name = scheduler_mapping.get(scheduler, scheduler) if scheduler else None
# Build output lines
lines = [prompt] if prompt else [""]
if negative_prompt:
metadata_parts.append(f"Negative prompt: {negative_prompt}")
lines.append(f"Negative prompt: {negative_prompt}")
# Format the second part (generation parameters)
params = []
# Add standard parameters in the correct order
if "steps" in metadata_dict:
add_param_if_not_none(params, "Steps", metadata_dict.get("steps"))
# Combine sampler and scheduler information
sampler_name = None
scheduler_name = None
if "sampler" in metadata_dict:
sampler = metadata_dict.get("sampler")
# Convert ComfyUI sampler names to user-friendly names
sampler_mapping = {
"euler": "Euler",
"euler_ancestral": "Euler a",
"dpm_2": "DPM2",
"dpm_2_ancestral": "DPM2 a",
"heun": "Heun",
"dpm_fast": "DPM fast",
"dpm_adaptive": "DPM adaptive",
"lms": "LMS",
"dpmpp_2s_ancestral": "DPM++ 2S a",
"dpmpp_sde": "DPM++ SDE",
"dpmpp_sde_gpu": "DPM++ SDE",
"dpmpp_2m": "DPM++ 2M",
"dpmpp_2m_sde": "DPM++ 2M SDE",
"dpmpp_2m_sde_gpu": "DPM++ 2M SDE",
"ddim": "DDIM",
}
sampler_name = sampler_mapping.get(sampler, sampler)
if "scheduler" in metadata_dict:
scheduler = metadata_dict.get("scheduler")
scheduler_mapping = {
"normal": "Simple",
"karras": "Karras",
"exponential": "Exponential",
"sgm_uniform": "SGM Uniform",
"sgm_quadratic": "SGM Quadratic",
}
scheduler_name = scheduler_mapping.get(scheduler, scheduler)
# Add combined sampler and scheduler information
params: list[str] = []
if steps is not None:
params.append(f"Steps: {steps}")
if sampler_name:
if scheduler_name:
params.append(f"Sampler: {sampler_name} {scheduler_name}")
else:
params.append(f"Sampler: {sampler_name}")
if cfg is not None:
params.append(f"CFG scale: {cfg}")
if seed is not None:
params.append(f"Seed: {seed}")
if size:
params.append(f"Size: {size}")
if clip_skip is not None:
try:
params.append(f"Clip skip: {abs(int(clip_skip))}")
except (ValueError, TypeError):
pass
additional_data = metadata_dict.get("additional_data", "")
if additional_data:
params.append(additional_data)
if ckpt_hash:
params.append(f"Model hash: {ckpt_hash[:10].upper()}")
if ckpt_display_name:
params.append(f"Model: {ckpt_display_name}")
if hashes:
params.append(f"Hashes: {json.dumps(hashes, separators=(',', ':'))}")
params.append("Version: ComfyUI")
if civitai_resources:
params.append(
f"Civitai resources: {json.dumps(civitai_resources, separators=(',', ':'))}"
)
# CFG scale (Use guidance if available, otherwise fall back to cfg_scale or cfg)
if "guidance" in metadata_dict:
add_param_if_not_none(params, "CFG scale", metadata_dict.get("guidance"))
elif "cfg_scale" in metadata_dict:
add_param_if_not_none(params, "CFG scale", metadata_dict.get("cfg_scale"))
elif "cfg" in metadata_dict:
add_param_if_not_none(params, "CFG scale", metadata_dict.get("cfg"))
# Seed
if "seed" in metadata_dict:
add_param_if_not_none(params, "Seed", metadata_dict.get("seed"))
# Size
if "size" in metadata_dict:
add_param_if_not_none(params, "Size", metadata_dict.get("size"))
# Model info
if "checkpoint" in metadata_dict:
# Ensure checkpoint is a string before processing
checkpoint = metadata_dict.get("checkpoint")
if checkpoint is not None:
# Get model hash
model_hash = self.get_checkpoint_hash(checkpoint)
# Extract basename without path
checkpoint_name = os.path.basename(checkpoint)
# Remove extension if present
checkpoint_name = os.path.splitext(checkpoint_name)[0]
# Add model hash if available
if model_hash:
params.append(
f"Model hash: {model_hash[:10]}, Model: {checkpoint_name}"
)
else:
params.append(f"Model: {checkpoint_name}")
# Add LoRA hashes if available
if lora_hashes:
lora_hash_parts = []
for lora_name, hash_value in lora_hashes.items():
lora_hash_parts.append(f"{lora_name}: {hash_value[:10]}")
if lora_hash_parts:
params.append(f'Lora hashes: "{", ".join(lora_hash_parts)}"')
# Combine all parameters with commas
metadata_parts.append(", ".join(params))
# Join all parts with a new line
return "\n".join(metadata_parts)
lines.append(", ".join(params))
return "\n".join(lines)
# credit to nkchocoai
# Add format_filename method to handle pattern substitution
@@ -298,7 +512,12 @@ class SaveImageLM:
key = parts[0]
if key == "seed" and "seed" in metadata_dict:
filename = filename.replace(segment, str(metadata_dict.get("seed", "")))
seed_value = metadata_dict.get("seed")
if seed_value is not None:
filename = filename.replace(segment, str(seed_value))
else:
# Fallback if seed was not captured by metadata collector
filename = filename.replace(segment, "0")
elif key == "width" and "size" in metadata_dict:
size = metadata_dict.get("size", "x")
w = size.split("x")[0] if isinstance(size, str) else size[0]
@@ -568,6 +787,8 @@ class SaveImageLM:
extra_pnginfo=None,
lossless_webp=True,
quality=100,
webp_method=6,
jpeg_subsampling=0,
embed_workflow=False,
save_with_metadata=True,
add_counter_to_filename=True,
@@ -603,7 +824,7 @@ class SaveImageLM:
img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
# Generate filename with counter if needed
base_filename = filename
base_filename = filename.replace("%batch_num%", str(i))
if add_counter_to_filename:
# Use counter + i to ensure unique filenames for all images in batch
current_counter = counter + i
@@ -622,15 +843,14 @@ class SaveImageLM:
elif file_format == "jpeg":
file = base_filename + ".jpg"
file_extension = ".jpg"
save_kwargs = {"quality": quality, "optimize": True}
save_kwargs = {"quality": quality, "optimize": True, "subsampling": jpeg_subsampling}
elif file_format == "webp":
file = base_filename + ".webp"
file_extension = ".webp"
# Add optimization param to control performance
save_kwargs = {
"quality": quality,
"lossless": lossless_webp,
"method": 0,
"method": webp_method,
}
else:
raise ValueError(f"Unsupported file format: {file_format}")
@@ -717,6 +937,8 @@ class SaveImageLM:
extra_pnginfo=None,
lossless_webp=True,
quality=100,
webp_method=6,
jpeg_subsampling=0,
embed_workflow=False,
save_with_metadata=True,
add_counter_to_filename=True,
@@ -746,6 +968,8 @@ class SaveImageLM:
extra_pnginfo,
lossless_webp,
quality,
webp_method,
jpeg_subsampling,
embed_workflow,
save_with_metadata,
add_counter_to_filename,
+20
View File
@@ -36,6 +36,7 @@ any_type = AnyType("*")
# Common methods extracted from lora_loader.py and lora_stacker.py
import os
import re
import logging
import copy
import sys
@@ -69,6 +70,25 @@ def extract_lora_name(lora_path):
return apply_lora_syntax_format(name_no_ext)
def parse_lora_syntax(text: str) -> list[dict]:
"""Parse <lora:name:strength> syntax from text input into a list of dicts.
Each entry contains: name, model_strength, clip_strength.
Supports both ``<lora:name:strength>`` and ``<lora:name:model_strength:clip_strength>``.
"""
pattern = r"<lora:([^:>]+):([^:>]+)(?::([^:>]+))?>"
matches = re.findall(pattern, text, re.IGNORECASE)
loras = []
for match in matches:
model_strength = float(match[1])
loras.append({
"name": match[0],
"model_strength": model_strength,
"clip_strength": float(match[2]) if match[2] else model_strength,
})
return loras
def get_loras_list(kwargs):
"""Helper to extract loras list from either old or new kwargs format"""
if "loras" not in kwargs:
+32 -15
View File
@@ -123,24 +123,39 @@ class AutomaticMetadataParser(RecipeMetadataParser):
if model_hash_from_hashes:
metadata["model_hash"] = model_hash_from_hashes
# Extract Lora hashes in alternative format
# Extract Lora hashes in alternative format.
# Run unconditionally (not just as fallback) so that
# non-empty hashes from Lora hashes fill in the gaps left
# by empty values in the Hashes JSON dict. Some WebUI
# builds write real hash values only to Lora hashes and
# leave the Hashes JSON values empty.
lora_hashes_match = re.search(self.LORA_HASHES_REGEX, params_section)
if not hashes_match and lora_hashes_match:
if lora_hashes_match:
try:
lora_hashes_str = lora_hashes_match.group(1)
lora_hash_entries = lora_hashes_str.split(', ')
# Initialize hashes dict if it doesn't exist
if "hashes" not in metadata:
metadata["hashes"] = {}
# Parse each lora hash entry (format: "name: hash")
for entry in lora_hash_entries:
if ': ' in entry:
lora_name, lora_hash = entry.split(': ', 1)
# Add as lora type in the same format as regular hashes
metadata["hashes"][f"lora:{lora_name}"] = lora_hash.strip()
lora_hash = lora_hash.strip()
if not lora_hash:
# Skip entries without a hash value
continue
# Initialize hashes dict if it doesn't exist
if "hashes" not in metadata:
metadata["hashes"] = {}
# Add as lora type in the same format as
# regular hashes. Only override an
# existing entry if its value is empty
# (Lora hashes is the more reliable
# source when Hashes JSON has blanks).
key = f"lora:{lora_name}"
existing = metadata["hashes"].get(key, "")
if not existing:
metadata["hashes"][key] = lora_hash
# Remove lora hashes from params section
params_section = params_section.replace(lora_hashes_match.group(0), '')
except Exception as e:
@@ -362,6 +377,12 @@ class AutomaticMetadataParser(RecipeMetadataParser):
# Only process lora or hypernet types
if not hash_key.startswith(("lora:", "hypernet:")):
continue
# Skip entries without a hash value — they can't be
# resolved via CivitAI and would only produce a
# useless "Deleted" entry in the recipe.
if not lora_hash:
continue
lora_type, lora_name = hash_key.split(':', 1)
@@ -387,11 +408,7 @@ class AutomaticMetadataParser(RecipeMetadataParser):
# Try to get info from Civitai
if metadata_provider:
try:
if lora_hash:
# If we have hash, use it for lookup
civitai_info = await metadata_provider.get_model_by_hash(lora_hash)
else:
civitai_info = None
civitai_info = await metadata_provider.get_model_by_hash(lora_hash)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
+53 -5
View File
@@ -514,11 +514,21 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
result["loras"].append(lora_entry)
# Process modelVersionIds from Civitai image API
# These are model version IDs returned at root level when meta doesn't contain resources
if "modelVersionIds" in metadata and isinstance(
metadata["modelVersionIds"], list
# Process modelVersionIds from Civitai image API.
# These are version IDs returned at root level of the API response.
# When resources or civitaiResources are already present in metadata
# (which they are when ?withMeta=true is passed), those sections have
# complete hash/type information — modelVersionIds is a fallback for
# when meta is null and only the flat ID list is available. Skipping
# it here avoids duplicates: the same file hash often resolves to
# different version IDs via hash lookup (resources) vs the original
# version ID in modelVersionIds, and both paths would create entries.
if (
"modelVersionIds" in metadata
and isinstance(metadata["modelVersionIds"], list)
and not result.get("loras")
):
for version_id in metadata["modelVersionIds"]:
version_id_str = str(version_id)
@@ -526,6 +536,13 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
if version_id_str in added_loras:
continue
# Skip if this version ID is already the recipe's checkpoint
# (resolved earlier from embedded resources/Model hash,
# avoiding a duplicate CivitAI API call).
existing_model = result.get("model")
if existing_model and str(existing_model.get("id")) == version_id_str:
continue
# Initialize lora entry with version ID
lora_entry = {
"id": version_id,
@@ -559,9 +576,40 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
)
if populated_entry is None:
continue # Skip invalid LoRA types
# Not a LoRA — try as checkpoint (only if we
# don't already have one). Reuses the same
# civitai_info from the API call above so no
# extra query is made.
if result["model"] is None:
checkpoint_entry = {
"id": version_id,
"modelId": 0,
"name": "Unknown Model",
"version": "",
"type": "checkpoint",
"existsLocally": False,
"localPath": None,
"file_name": "",
"hash": "",
"thumbnailUrl": (
"/loras_static/images/no-preview.png"
),
"baseModel": "",
"size": 0,
"downloadUrl": "",
"isDeleted": False,
}
cp_populated = await (
self.populate_checkpoint_from_civitai(
checkpoint_entry, civitai_info
)
)
if cp_populated.get("modelId"):
result["model"] = cp_populated
continue # Not a LoRA, don't add to loras
lora_entry = populated_entry
except Exception as e:
logger.error(
f"Error fetching Civitai info for model version {version_id}: {e}"
+165
View File
@@ -0,0 +1,165 @@
"""HTTP route handlers for agent skill endpoints.
These handlers expose the :class:`AgentService` via HTTP, allowing the
frontend to list available skills and execute them on selected models.
Progress is reported via WebSocket broadcast.
"""
from __future__ import annotations
import asyncio
import logging
from typing import Any, Dict
from aiohttp import web
from ...services.agent import AgentService, AgentProgressReporter
from ...services.llm_service import LLMNotConfiguredError
logger = logging.getLogger(__name__)
class AgentHandler:
"""HTTP handler for agent skill operations."""
def __init__(self, agent_service: AgentService | None = None) -> None:
self._agent_service = agent_service
async def _ensure_service(self) -> AgentService:
if self._agent_service is None:
self._agent_service = await AgentService.get_instance()
return self._agent_service
# ------------------------------------------------------------------
# GET /api/lm/agent/skills
# ------------------------------------------------------------------
async def get_agent_skills(self, request: web.Request) -> web.Response:
"""Return a list of available agent skills."""
service = await self._ensure_service()
skills = await service.list_skills()
return web.json_response({"skills": skills})
# ------------------------------------------------------------------
# POST /api/lm/agent/execute/{skill_name}
# ------------------------------------------------------------------
async def execute_agent_skill(self, request: web.Request) -> web.Response:
"""Execute an agent skill on the provided model paths.
Request body::
{"model_paths": ["/path/to/model1.safetensors", ...], "options": {}}
Returns immediately with a task ID. Execution runs in the
background; progress and completion are pushed via WebSocket
events of type ``agent_progress``.
"""
skill_name = request.match_info.get("skill_name", "")
if not skill_name:
return web.json_response(
{"error": "Skill name is required"}, status=400
)
try:
body = await request.json()
except Exception:
return web.json_response(
{"error": "Invalid JSON body"}, status=400
)
model_paths = body.get("model_paths", [])
if not model_paths or not isinstance(model_paths, list):
return web.json_response(
{"error": "model_paths must be a non-empty array"},
status=400,
)
service = await self._ensure_service()
# Validate LLM configuration early for skills that need it
# (fail fast rather than after starting background work)
try:
from ...services.llm_service import LLMService
llm = await LLMService.get_instance()
if not llm.is_configured():
return web.json_response(
{
"error": "LLM provider is not configured. "
"Enable it in Settings → AI Provider.",
},
status=400,
)
except Exception as exc:
logger.error("Failed to check LLM configuration: %s", exc)
# Launch execution in the background
progress_reporter = AgentProgressReporter()
logger.info(
"LLM enrichment '%s' starting for %d model(s)",
skill_name, len(model_paths),
)
async def _run() -> None:
try:
result = await service.execute_skill(
skill_name=skill_name,
input_data={"model_paths": model_paths},
progress_callback=progress_reporter,
)
logger.info(
"LLM enrichment '%s' finished: success=%s, summary='%s', errors=%s",
skill_name, result.success, result.summary, result.errors,
)
except LLMNotConfiguredError as exc:
logger.warning("LLM enrichment '%s' not configured: %s", skill_name, exc)
await progress_reporter.on_progress(
{
"type": "agent_progress",
"skill": skill_name,
"status": "error",
"error": str(exc),
}
)
except Exception as exc:
logger.error("LLM enrichment '%s' failed: %s", skill_name, exc, exc_info=True)
await progress_reporter.on_progress(
{
"type": "agent_progress",
"skill": skill_name,
"status": "error",
"error": str(exc),
}
)
# Fire and forget — progress comes via WebSocket
asyncio.create_task(_run())
return web.json_response(
{
"status": "started",
"skill": skill_name,
"model_count": len(model_paths),
}
)
# ------------------------------------------------------------------
# POST /api/lm/agent/cancel
# ------------------------------------------------------------------
async def cancel_agent_skill(self, request: web.Request) -> web.Response:
"""Cancel a running agent skill.
NOTE: Cancellation is a stub for now the AgentService processes
models sequentially and does not yet support mid-execution
cancellation. This endpoint exists for API completeness.
"""
# TODO: implement cooperative cancellation in AgentService
return web.json_response(
{"status": "acknowledged", "note": "Cancellation not yet implemented"},
status=200,
)
+508
View File
@@ -0,0 +1,508 @@
"""Handlers for Hugging Face model listing and download.
Minimal MVP implementation uses direct HTTP to the HF API for file
listing and the project's existing aiohttp-based Downloader for
downloading. No huggingface_hub dependency required.
"""
from __future__ import annotations
import json
import logging
import os
import re
from typing import Any
import aiohttp
from aiohttp import web
from ...config import config
from ...services.downloader import (
DownloadProgress,
get_downloader,
)
from ...services.aria2_downloader import Aria2Downloader
from ...services.settings_manager import get_settings_manager
from ...services.service_registry import ServiceRegistry
from ...services.websocket_manager import ws_manager
from ...utils.constants import MODEL_FILE_EXTENSIONS
from ...utils.metadata_manager import MetadataManager
from ...utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
logger = logging.getLogger(__name__)
_DEFAULT_MODEL_CLASS = LoraMetadata
_DEFAULT_SCANNER_GETTER = "get_lora_scanner"
# Shared aiohttp session for HF API calls (created on first use)
_hf_api_session: aiohttp.ClientSession | None = None
async def _get_hf_api_session() -> aiohttp.ClientSession:
"""Get or create the shared aiohttp session for HF API calls."""
global _hf_api_session # needed because we reassign the module-level name
if _hf_api_session is None or _hf_api_session.closed:
_hf_api_session = aiohttp.ClientSession(
headers={"User-Agent": "ComfyUI-LoRA-Manager/1.0"},
timeout=aiohttp.ClientTimeout(total=30),
)
return _hf_api_session
async def close_hf_api_session() -> None:
"""Close the shared HF API session, if it was ever created."""
global _hf_api_session
if _hf_api_session is not None and not _hf_api_session.closed:
await _hf_api_session.close()
_hf_api_session = None
def _infer_model_type(model_root: str) -> tuple[Any, str]:
"""Determine model class and scanner by matching ``model_root`` against the
configured root paths for each model type (from ``Config``).
The ``model_root`` value comes from the frontend's model-root dropdown,
which is populated from the current page's scanner roots. By checking
which scanner's root list it belongs to, we avoid fragile heuristics
like substring-matching path names.
"""
norm = os.path.normpath(model_root).replace(os.sep, "/")
# LoRA roots
for p in (config.loras_roots or []) + (config.extra_loras_roots or []):
if os.path.normpath(p).replace(os.sep, "/") == norm:
return LoraMetadata, "get_lora_scanner"
# Checkpoint / UNet roots
for p in (
(config.checkpoints_roots or [])
+ (config.extra_checkpoints_roots or [])
+ (config.unet_roots or [])
+ (config.extra_unet_roots or [])
):
if os.path.normpath(p).replace(os.sep, "/") == norm:
return CheckpointMetadata, "get_checkpoint_scanner"
# Embedding roots
for p in (config.embeddings_roots or []) + (config.extra_embeddings_roots or []):
if os.path.normpath(p).replace(os.sep, "/") == norm:
return EmbeddingMetadata, "get_embedding_scanner"
# Fallback — should not happen in normal use
logger.warning(
"Could not determine model type for root '%s'; defaulting to LoRA",
model_root,
)
return _DEFAULT_MODEL_CLASS, _DEFAULT_SCANNER_GETTER
async def _save_hf_metadata(dest_path: str, repo: str, model_root: str) -> None:
"""Create a proper .metadata.json and add the model to the scanner cache.
Uses ``MetadataManager.create_default_metadata()`` which computes the
SHA256 hash, extracts safetensors header metadata (base_model), and
produces a fully-populated ``LoraMetadata`` (or ``CheckpointMetadata`` /
``EmbeddingMetadata``) object. We then overlay HF-specific fields and
register the model in the in-memory scanner cache so it appears
immediately without a full filesystem walk.
"""
try:
hf_url = f"https://huggingface.co/{repo}"
model_class, scanner_getter_name = _infer_model_type(model_root)
# 1. Create proper metadata (computes SHA256, reads safetensors headers)
metadata = await MetadataManager.create_default_metadata(
dest_path, model_class=model_class
)
if metadata is None:
logger.warning("create_default_metadata returned None for %s", dest_path)
return
# 2. Overlay HF-specific fields
metadata._unknown_fields["hf_url"] = hf_url
metadata.from_civitai = False # HF models are not from CivitAI
metadata_dict = metadata.to_dict()
if "trainedWords" in metadata_dict and not metadata_dict["trainedWords"]:
del metadata_dict["trainedWords"]
# 3. Save metadata atomically
await MetadataManager.save_metadata(dest_path, metadata_dict)
logger.info("Saved HF metadata (with hf_url) for %s", dest_path)
# 4. Determine relative folder path for cache
# model_root is an absolute path; dest_path is under it
folder = ""
if os.path.isabs(model_root) and dest_path.startswith(model_root):
rel = os.path.relpath(os.path.dirname(dest_path), model_root)
folder = rel.replace(os.sep, "/") if rel != "." else ""
# 5. Add to scanner cache (same as CivitAI's _execute_download does)
scanner_getter = getattr(ServiceRegistry, scanner_getter_name, None)
if scanner_getter is not None:
scanner = await scanner_getter()
if scanner is not None:
metadata_dict = metadata.to_dict()
metadata_dict["hf_url"] = hf_url
await scanner.add_model_to_cache(metadata_dict, folder)
logger.info("Added %s to scanner cache (folder=%s)", dest_path, folder)
except Exception as exc:
logger.warning("Failed to save HF metadata for %s: %s", dest_path, exc)
def _find_matching_root(dest_dir: str) -> str | None:
"""Walk up *dest_dir* to find which configured scanner root it belongs to."""
norm = os.path.normpath(dest_dir).replace(os.sep, "/")
all_roots = []
for root_list in (
config.loras_roots or [],
config.extra_loras_roots or [],
config.checkpoints_roots or [],
config.extra_checkpoints_roots or [],
config.unet_roots or [],
config.extra_unet_roots or [],
config.embeddings_roots or [],
config.extra_embeddings_roots or [],
):
all_roots.extend([os.path.normpath(p).replace(os.sep, "/") for p in root_list])
# Find the longest matching prefix
match: str | None = None
for root in all_roots:
if norm.startswith(root):
if match is None or len(root) > len(match):
match = root
return match
async def _add_to_scanner_cache(dest_path: str, metadata: dict[str, Any]) -> None:
model_dir = os.path.dirname(dest_path)
model_root = _find_matching_root(model_dir)
if not model_root:
raise ValueError(f"File path {dest_path} is not within any configured scanner root")
scanner_getter_name = _infer_model_type(model_root)[1]
scanner_getter = getattr(ServiceRegistry, scanner_getter_name, None)
if scanner_getter is None:
raise RuntimeError(f"Scanner getter '{scanner_getter_name}' not found in ServiceRegistry")
scanner = await scanner_getter()
if scanner is None:
raise RuntimeError(f"Scanner '{scanner_getter_name}' returned None")
await scanner.update_single_model_cache(dest_path, dest_path, metadata)
class HfHandler:
"""Handle Hugging Face model browsing and download."""
async def set_hf_url(self, request: web.Request) -> web.Response:
try:
payload: dict[str, Any] = await request.json()
except json.JSONDecodeError:
return web.json_response({"success": False, "error": "Invalid JSON"}, status=400)
file_path = (payload.get("file_path") or "").strip()
hf_url = (payload.get("hf_url") or "").strip()
if not file_path or not hf_url:
return web.json_response(
{"success": False, "error": "Missing required fields: 'file_path' and 'hf_url'"},
status=400,
)
m = re.match(r"^https?://huggingface\.co/([^/]+/[^/]+)/?$", hf_url)
if not m:
return web.json_response(
{
"success": False,
"error": "Invalid HuggingFace URL. Expected format: https://huggingface.co/user/repo",
},
status=400,
)
if not os.path.isfile(file_path):
return web.json_response(
{"success": False, "error": f"File not found: {file_path}"},
status=404,
)
model_root = _find_matching_root(os.path.dirname(file_path))
if not model_root:
return web.json_response(
{
"success": False,
"error": "File is not within any configured model directory. Cannot link to HuggingFace.",
},
status=400,
)
try:
existing = await MetadataManager.load_metadata_payload(file_path)
if existing.get("hf_url") == hf_url:
return web.json_response({
"success": True,
"message": "hf_url already set",
"hf_url": hf_url,
})
existing["hf_url"] = hf_url
existing["from_civitai"] = False
await MetadataManager.save_metadata(file_path, existing)
await _add_to_scanner_cache(file_path, existing)
logger.info("Set hf_url=%s for %s", hf_url, file_path)
return web.json_response({
"success": True,
"message": f"hf_url set to {hf_url}",
"hf_url": hf_url,
})
except Exception as exc:
logger.error("Failed to set hf_url for %s: %s", file_path, exc)
return web.json_response(
{"success": False, "error": str(exc)},
status=500,
)
async def get_hf_repo_files(self, request: web.Request) -> web.Response:
"""List model-weight files from a HF repo with real file sizes.
Uses the HF tree API endpoint which returns accurate file sizes
(including LFS-tracked files), unlike the model info endpoint.
"""
repo = request.query.get("repo", "").strip()
if not repo or "/" not in repo:
return web.json_response(
{"error": "Missing or invalid 'repo' parameter (expected user/repo)"},
status=400,
)
url = f"https://huggingface.co/api/models/{repo}/tree/main"
try:
session = await _get_hf_api_session()
async with session.get(url) as resp:
if resp.status == 404:
return web.json_response(
{"error": f"Repo '{repo}' not found"}, status=404
)
if resp.status != 200:
text = await resp.text()
return web.json_response(
{"error": f"HF API error {resp.status}: {text[:200]}"},
status=resp.status,
)
tree: list[dict[str, Any]] = await resp.json()
except Exception as exc:
logger.error("Failed to fetch HF repo files: %s", exc)
return web.json_response({"error": str(exc)}, status=502)
files: list[dict[str, Any]] = []
for entry in tree:
path: str = entry.get("path", "")
ext = os.path.splitext(path)[1].lower()
if ext not in MODEL_FILE_EXTENSIONS:
continue
size = entry.get("size", 0) or 0
if size == 0 and "lfs" in entry:
size = entry["lfs"].get("size", 0) or 0
files.append({
"filename": path,
"size": size,
})
files.sort(key=lambda f: f["size"], reverse=True)
return web.json_response(files)
async def download_hf_model(self, request: web.Request) -> web.Response:
"""Download a single file from Hugging Face into the model directory.
POST JSON body::
{
"repo": "dx8152/Flux2-Klein-9B-Consistency",
"filename": "Flux2-Klein-9B-consistency-V2.safetensors",
"revision": "main",
"model_root": "loras",
"relative_path": "",
"use_default_paths": false,
"download_id": "optional-batch-id"
}
If ``download_id`` is provided, real-time progress (bytes, speed,
percentage) is broadcast via the WebSocket progress system, matching
the CivitAI download experience.
Respects the ``download_backend`` setting (``aria2`` or ``default``).
"""
try:
payload: dict[str, Any] = await request.json()
except json.JSONDecodeError:
return web.json_response({"error": "Invalid JSON"}, status=400)
repo = (payload.get("repo") or "").strip()
filename = (payload.get("filename") or "").strip()
revision = (payload.get("revision") or "main").strip()
model_root = (payload.get("model_root") or "").strip()
relative_path = (payload.get("relative_path") or "").strip()
use_default_paths = bool(payload.get("use_default_paths", False))
download_id: str | None = payload.get("download_id")
logger.info(
"download_hf_model: repo=%s file=%s root=%s download_id=%s",
repo, filename, model_root, download_id,
)
if not repo or not filename:
return web.json_response(
{"error": "Missing required fields: 'repo' and 'filename'"}, status=400
)
# Validate repo format — must be user/repo_name
if repo.count("/") != 1 or not re.match(r"^[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+$", repo):
return web.json_response({"error": f"Invalid repo format: {repo}"}, status=400)
author, repo_name = repo.split("/", 1)
if ".." in (author, repo_name) or "." in (author, repo_name):
return web.json_response({"error": f"Invalid repo format: {repo}"}, status=400)
# Validate filename — must not contain path traversal
if ".." in filename:
return web.json_response({"error": "Invalid filename"}, status=400)
# Validate relative_path — must not be absolute or escape base directory
if relative_path:
if os.path.isabs(relative_path):
return web.json_response({"error": "relative_path must not be absolute"}, status=400)
if ".." in relative_path.split("/") or "\\" in relative_path:
return web.json_response({"error": "Invalid relative_path"}, status=400)
# Use model_root directly as the base directory — same approach as
# CivitAI's download path (download_manager.py). No realpath, no
# allowed-roots validation, no path-traversal check; those are
# unnecessary when the frontend sends the path from its own dropdown
# (populated from scanner roots). Using the "business path" directly
# keeps dest_path consistent with scanner roots so that later folder
# derivation (in _save_hf_metadata) works correctly.
if os.path.isabs(model_root):
base_dir = os.path.normpath(model_root)
else:
base_dir = os.path.normpath(os.path.join(os.getcwd(), "models", model_root))
if use_default_paths:
target_dir = os.path.join(base_dir, "huggingface", author, repo_name)
elif relative_path:
target_dir = os.path.join(base_dir, relative_path)
else:
target_dir = base_dir
# Strip HF repo subdirectory — "diffusion_models/xxx.safetensors"
# is an HF repo convention, not meaningful for local storage.
file_base = os.path.basename(filename)
os.makedirs(target_dir, exist_ok=True)
dest_path = os.path.join(target_dir, file_base)
# Check if already exists (simple skip)
if os.path.exists(dest_path) and os.path.getsize(dest_path) > 0:
logger.info("download_hf_model: file already exists, skipping — %s", dest_path)
return web.json_response({
"success": True,
"message": f"File already exists: {dest_path}",
"path": dest_path,
})
# Build HF resolve URL
resolve_url = (
f"https://huggingface.co/{repo}/resolve/{revision}/{filename}"
)
# Set up progress callback if download_id is provided
progress_callback = None
if download_id:
async def _progress_callback(
progress: float | DownloadProgress,
snapshot: DownloadProgress | None = None,
) -> None:
percent = 0.0
metrics = snapshot if isinstance(snapshot, DownloadProgress) else None
if isinstance(progress, DownloadProgress):
percent = progress.percent_complete
metrics = progress
elif isinstance(snapshot, DownloadProgress):
percent = snapshot.percent_complete
else:
percent = float(progress)
broadcast: dict[str, Any] = {
"status": "progress",
"progress": round(percent),
}
if metrics:
broadcast["bytes_downloaded"] = metrics.bytes_downloaded
broadcast["total_bytes"] = metrics.total_bytes
broadcast["bytes_per_second"] = metrics.bytes_per_second
await ws_manager.broadcast_download_progress(download_id, broadcast)
progress_callback = _progress_callback
# Respect download backend setting (aria2 vs default)
download_backend = (
get_settings_manager().get("download_backend", "default")
)
if download_backend == "aria2":
aria2 = await Aria2Downloader.get_instance()
aid = download_id or f"hf_{repo}_{filename}"
try:
hf_success, hf_result = await aria2.download_file(
url=resolve_url,
save_path=dest_path,
download_id=aid,
progress_callback=progress_callback,
)
if hf_success:
await _save_hf_metadata(dest_path, repo, model_root)
return web.json_response({
"success": True,
"message": f"Downloaded to {dest_path}",
"path": dest_path,
})
else:
return web.json_response(
{"success": False, "error": hf_result or "aria2 download failed"},
status=500,
)
except Exception as exc:
logger.error("HF download (aria2) failed: %s", exc)
return web.json_response(
{"success": False, "error": str(exc)}, status=500
)
# Default: use built-in aiohttp Downloader
downloader = await get_downloader()
try:
success, result = await downloader.download_file(
url=resolve_url,
save_path=dest_path,
use_auth=False,
allow_resume=True,
progress_callback=progress_callback,
)
if success:
await _save_hf_metadata(dest_path, repo, model_root)
return web.json_response({
"success": True,
"message": f"Downloaded to {result}",
"path": result,
})
else:
return web.json_response(
{"success": False, "error": result or "Download failed"},
status=500,
)
except Exception as exc:
logger.error("HF download failed: %s", exc)
return web.json_response(
{"success": False, "error": str(exc)}, status=500
)
+627 -99
View File
@@ -38,6 +38,12 @@ from ...services.settings_manager import get_settings_manager
from ...services.websocket_manager import ws_manager
from ...services.downloader import get_downloader
from ...services.errors import ResourceNotFoundError
from ...services.llm_service import (
PROVIDER_PRESETS,
fetch_ollama_models,
get_all_provider_models,
get_provider_model_ids,
)
from ...services.cache_health_monitor import CacheHealthMonitor, CacheHealthStatus
from ...utils.models import BaseModelMetadata
from ...utils.constants import (
@@ -48,8 +54,13 @@ from ...utils.constants import (
SUPPORTED_MEDIA_EXTENSIONS,
VALID_LORA_TYPES,
)
from .hf_handlers import HfHandler
from .agent_handlers import AgentHandler
from ...utils.civitai_utils import rewrite_preview_url
from ...utils.example_images_paths import is_valid_example_images_root
from ...utils.example_images_paths import (
find_non_compliant_items_in_example_images_root,
is_valid_example_images_root,
)
from ...utils.lora_metadata import extract_trained_words
from ...utils.session_logging import get_standalone_session_log_snapshot
from ...utils.usage_stats import UsageStats
@@ -411,9 +422,10 @@ class PromptServerProtocol(Protocol):
"""Subset of PromptServer used by the handlers."""
instance: "PromptServerProtocol"
sockets: dict # maps clientId (sid) → WebSocketResponse
def send_sync(
self, event: str, payload: dict
self, event: str, payload: dict | None = None, sid: str | None = None
) -> None: # pragma: no cover - protocol
...
@@ -468,89 +480,167 @@ class BackupServiceProtocol(Protocol):
class NodeRegistry:
"""Thread-safe registry for tracking LoRA nodes in active workflows."""
"""Thread-safe registry for tracking LoRA nodes across ComfyUI tabs.
Each connected ComfyUI browser tab (identified by its ``sid`` / ``clientId``)
registers its own set of workflow nodes. Queries merge all known tabs into
a single result so that the calling LM panel always sees *every* available
target node, regardless of which tab responded fastest.
"""
def __init__(self) -> None:
self._lock = asyncio.Lock()
self._nodes: Dict[str, dict] = {}
self._registry_updated = asyncio.Event()
# sid → {unique_id → node_info}
self._tab_nodes: Dict[str, Dict[str, dict]] = {}
self._ready = asyncio.Event()
self._waiting_clients: set[str] = set()
@property
def pending_client_count(self) -> int:
"""Number of clients that have not yet responded in the current refresh cycle."""
return len(self._waiting_clients)
# ------------------------------------------------------------------
# Helpers to build one node dict (extracted so it's reused for each tab)
# ------------------------------------------------------------------
@staticmethod
def _build_node_dict(node: dict) -> dict:
node_id = node["node_id"]
graph_id = str(node["graph_id"])
unique_id = f"{graph_id}:{node_id}"
node_type = node.get("type", "")
type_id = NODE_TYPES.get(node_type, 0)
bgcolor = node.get("bgcolor") or DEFAULT_NODE_COLOR
raw_capabilities = node.get("capabilities")
capabilities: dict = {}
if isinstance(raw_capabilities, dict):
capabilities = dict(raw_capabilities)
raw_widget_names: list | None = node.get("widget_names")
if not isinstance(raw_widget_names, list):
capability_widget_names = capabilities.get("widget_names")
raw_widget_names = (
capability_widget_names
if isinstance(capability_widget_names, list)
else None
)
widget_names: list[str] = []
if isinstance(raw_widget_names, list):
widget_names = [
str(widget_name)
for widget_name in raw_widget_names
if isinstance(widget_name, str) and widget_name
]
if widget_names:
capabilities["widget_names"] = widget_names
else:
capabilities.pop("widget_names", None)
if "supports_lora" in capabilities:
capabilities["supports_lora"] = bool(capabilities["supports_lora"])
comfy_class = node.get("comfy_class")
if not isinstance(comfy_class, str) or not comfy_class:
comfy_class = node_type if isinstance(node_type, str) else None
return {
"id": node_id,
"graph_id": graph_id,
"graph_name": node.get("graph_name"),
"unique_id": unique_id,
"bgcolor": bgcolor,
"title": node.get("title"),
"type": type_id,
"type_name": node_type,
"comfy_class": comfy_class,
"capabilities": capabilities,
"widget_names": widget_names,
"mode": node.get("mode"),
"marker_role": node.get("marker_role"),
}
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
async def register_nodes(self, sid: str, nodes: list[dict]) -> None:
"""Register/replace the node list for a single ComfyUI tab (identified by *sid*)."""
tab_nodes: dict[str, dict] = {}
for node in nodes:
nd = self._build_node_dict(node)
tab_nodes[nd["unique_id"]] = nd
async def register_nodes(self, nodes: list[dict]) -> None:
async with self._lock:
self._nodes.clear()
for node in nodes:
node_id = node["node_id"]
graph_id = str(node["graph_id"])
unique_id = f"{graph_id}:{node_id}"
node_type = node.get("type", "")
type_id = NODE_TYPES.get(node_type, 0)
bgcolor = node.get("bgcolor") or DEFAULT_NODE_COLOR
raw_capabilities = node.get("capabilities")
capabilities: dict = {}
if isinstance(raw_capabilities, dict):
capabilities = dict(raw_capabilities)
prev_count = len(self._tab_nodes.get(sid, {}))
self._tab_nodes[sid] = tab_nodes
self._waiting_clients.discard(sid)
if not self._waiting_clients:
self._ready.set()
total_tabs = len(self._tab_nodes)
raw_widget_names: list | None = node.get("widget_names")
if not isinstance(raw_widget_names, list):
capability_widget_names = capabilities.get("widget_names")
raw_widget_names = (
capability_widget_names
if isinstance(capability_widget_names, list)
else None
)
if len(nodes) != prev_count or len(nodes) > 0:
logger.debug(
"[LM:Registry] stored %s nodes (was %s) for client %s (total tabs: %s)",
len(nodes), prev_count, sid, total_tabs,
)
widget_names: list[str] = []
if isinstance(raw_widget_names, list):
widget_names = [
str(widget_name)
for widget_name in raw_widget_names
if isinstance(widget_name, str) and widget_name
]
def prepare_for_refresh(self, active_sids: list[str]) -> None:
"""Set the list of client IDs we expect to hear from during the next refresh cycle."""
self._ready.clear()
self._waiting_clients = set(active_sids)
if widget_names:
capabilities["widget_names"] = widget_names
else:
capabilities.pop("widget_names", None)
if "supports_lora" in capabilities:
capabilities["supports_lora"] = bool(capabilities["supports_lora"])
comfy_class = node.get("comfy_class")
if not isinstance(comfy_class, str) or not comfy_class:
comfy_class = node_type if isinstance(node_type, str) else None
self._nodes[unique_id] = {
"id": node_id,
"graph_id": graph_id,
"graph_name": node.get("graph_name"),
"unique_id": unique_id,
"bgcolor": bgcolor,
"title": node.get("title"),
"type": type_id,
"type_name": node_type,
"comfy_class": comfy_class,
"capabilities": capabilities,
"widget_names": widget_names,
"mode": node.get("mode"),
}
logger.debug("Registered %s nodes in registry", len(nodes))
self._registry_updated.set()
async def get_registry(self) -> dict:
async with self._lock:
return {
"nodes": dict(self._nodes),
"node_count": len(self._nodes),
}
async def wait_for_update(self, timeout: float = 1.0) -> bool:
self._registry_updated.clear()
async def wait_for_all(self, timeout: float = 2.0) -> bool:
"""Block until every client in the current waiting set has responded
(or *timeout* seconds elapse). Returns ``True`` if all responded."""
if not self._waiting_clients:
return True
try:
await asyncio.wait_for(self._registry_updated.wait(), timeout=timeout)
await asyncio.wait_for(self._ready.wait(), timeout=timeout)
return True
except asyncio.TimeoutError:
return False
async def get_merged_registry(self, active_sids: set[str] | None = None) -> dict:
"""Return the union of all known tab nodes, pruning any tab that is no
longer connected."""
async with self._lock:
# Garbage-collect stale entries (disconnected tabs)
stale_sids = []
if active_sids is not None:
for sid in list(self._tab_nodes):
if sid not in active_sids:
stale_sids.append(sid)
del self._tab_nodes[sid]
if stale_sids:
logger.debug(
"[LM:Registry] GC pruned %s disconnected tabs: %s",
len(stale_sids), stale_sids,
)
merged: dict[str, dict] = {}
tab_info: dict[str, dict] = {}
for sid, nodes in self._tab_nodes.items():
tab_info[sid] = {
"node_count": len(nodes),
"graph_names": list(
{
n.get("graph_name")
for n in nodes.values()
if n.get("graph_name")
}
),
}
merged.update(nodes)
return {
"nodes": merged,
"node_count": len(merged),
"tab_count": len(self._tab_nodes),
"tabs": tab_info,
}
class HealthCheckHandler:
async def health_check(self, request: web.Request) -> web.Response:
@@ -1328,6 +1418,10 @@ class SettingsHandler:
"folder_paths",
"libraries",
"active_library",
# Sensitive — never expose the actual value to the frontend;
# frontend receives a boolean instead (*_set).
"civitai_api_key",
"llm_api_key",
}
)
@@ -1382,6 +1476,11 @@ class SettingsHandler:
value = self._settings.get(key)
if value is not None:
response_data[key] = value
# Sensitive fields: only expose a boolean indicating whether set
raw_key = self._settings.get("civitai_api_key")
response_data["civitai_api_key_set"] = bool(raw_key)
raw_llm_key = self._settings.get("llm_api_key")
response_data["llm_api_key_set"] = bool(raw_llm_key)
settings_file = getattr(self._settings, "settings_file", None)
if settings_file:
response_data["settings_file"] = settings_file
@@ -1463,6 +1562,11 @@ class SettingsHandler:
{"success": False, "error": validation_error}
)
if key == "update_channel" and value not in ("release", "nightly"):
return web.json_response(
{"success": False, "error": "update_channel must be 'release' or 'nightly'"}
)
if value == "__DELETE__" and key in (
"proxy_username",
"proxy_password",
@@ -1471,7 +1575,11 @@ class SettingsHandler:
else:
self._settings.set(key, value)
if key == "enable_metadata_archive_db":
if key in (
"enable_metadata_archive_db",
"enable_civarchive_api",
"metadata_provider_order",
):
await self._metadata_provider_updater()
if key in self._PROXY_KEYS:
@@ -1486,18 +1594,78 @@ class SettingsHandler:
logger.error("Error updating settings: %s", exc, exc_info=True)
return web.Response(status=500, text=str(exc))
async def get_llm_models(self, request: web.Request) -> web.Response:
"""Return the model list for a provider.
For ``ollama`` the list is fetched live from the local Ollama API
(only models actually pulled locally are shown). For all other
providers the opencode model catalog is used.
Query parameters:
provider (required): Internal provider id (``openai``, ``ollama``, etc.).
Returns:
``{"success": true, "models": ["gpt-4o", ...]}``.
"""
provider_id = request.query.get("provider", "").strip()
if not provider_id:
return web.json_response(
{"success": False, "error": "provider query parameter is required", "models": []},
status=400,
)
try:
if provider_id == "ollama":
api_base = request.query.get("api_base", "").strip() or self._settings.get("llm_api_base", "")
if not api_base:
api_base = "http://localhost:11434/v1"
models = await fetch_ollama_models(api_base)
else:
models = await get_provider_model_ids(provider_id)
return web.json_response({"success": True, "models": models})
except Exception as exc:
logger.warning("get_llm_models failed for %s: %s", provider_id, exc)
return web.json_response(
{"success": False, "error": str(exc), "models": []},
status=500,
)
def _validate_example_images_path(self, folder_path: str) -> str | None:
if not os.path.exists(folder_path):
return f"Path does not exist: {folder_path}"
if not os.path.isdir(folder_path):
return "Please set a dedicated folder for example images."
if not self._is_dedicated_example_images_folder(folder_path):
offending = find_non_compliant_items_in_example_images_root(folder_path)
if offending:
items_str = ", ".join(repr(item) for item in offending[:5])
if len(offending) > 5:
items_str += f" … and {len(offending) - 5} more"
return (
f"The folder contains items that are not valid example image "
f"folders: {items_str}. Please use a dedicated, empty folder "
f"for example images to prevent accidental data loss."
)
return "Please set a dedicated folder for example images."
return None
def _is_dedicated_example_images_folder(self, folder_path: str) -> bool:
return is_valid_example_images_root(folder_path)
async def get_provider_models(self, request: web.Request) -> web.Response:
"""Return the model catalog for all preset providers.
This endpoint is called asynchronously by the settings UI so that
page rendering never blocks on the remote model catalog fetch.
"""
catalog_provider_ids = [p for p in PROVIDER_PRESETS if p != "custom"]
try:
provider_models = await get_all_provider_models(catalog_provider_ids)
return web.json_response({"success": True, "models": provider_models})
except Exception as exc:
logger.warning("Failed to fetch provider models: %s", exc)
return web.json_response({"success": False, "models": {}, "error": str(exc)})
class UsageStatsHandler:
def __init__(self, usage_stats_factory: UsageStatsFactory = UsageStats) -> None:
@@ -1625,6 +1793,124 @@ class LoraCodeHandler:
logger.error("Failed to update lora code: %s", exc, exc_info=True)
return web.json_response({"success": False, "error": str(exc)}, status=500)
async def get_update_lora_code(self, request: web.Request) -> web.Response:
"""GET version of update_lora_code — reads parameters from query string.
Query params:
lora_code (required) the LoRA syntax to send
mode (optional) "append" (default) or "replace"
node_id (repeatable) target node id(s), e.g. node_id=3&node_id=5
node_ids (optional) JSON-encoded array for complex references with graph_id:
[{"node_id":3,"graph_id":"g1"}, ...]
"""
try:
node_ids_raw = request.query.get("node_ids")
node_id_list = request.query.getall("node_id", [])
lora_code = request.query.get("lora_code", "")
mode = request.query.get("mode", "append")
if not lora_code:
return web.json_response(
{"success": False, "error": "Missing lora_code parameter"},
status=400,
)
node_ids = None
if node_ids_raw:
try:
node_ids = json.loads(node_ids_raw)
except (json.JSONDecodeError, TypeError):
return web.json_response(
{"success": False, "error": "node_ids must be a valid JSON array"},
status=400,
)
if not isinstance(node_ids, list) or not node_ids:
return web.json_response(
{"success": False, "error": "node_ids must be a non-empty JSON array"},
status=400,
)
elif node_id_list:
node_ids = node_id_list
results = []
if node_ids is None:
try:
self._prompt_server.instance.send_sync(
"lora_code_update",
{"id": -1, "lora_code": lora_code, "mode": mode},
)
results.append({"node_id": "broadcast", "success": True})
except Exception as exc: # pragma: no cover - defensive logging
logger.error("Error broadcasting lora code: %s", exc)
results.append(
{"node_id": "broadcast", "success": False, "error": str(exc)}
)
else:
for entry in node_ids:
node_identifier = entry
graph_identifier = None
if isinstance(entry, dict):
node_identifier = entry.get("node_id")
graph_identifier = entry.get("graph_id")
if node_identifier is None:
results.append(
{
"node_id": node_identifier,
"graph_id": graph_identifier,
"success": False,
"error": "Missing node_id parameter",
}
)
continue
try:
parsed_node_id = int(node_identifier)
except (TypeError, ValueError):
parsed_node_id = node_identifier
payload = {
"id": parsed_node_id,
"lora_code": lora_code,
"mode": mode,
}
if graph_identifier is not None:
payload["graph_id"] = str(graph_identifier)
try:
self._prompt_server.instance.send_sync(
"lora_code_update",
payload,
)
results.append(
{
"node_id": parsed_node_id,
"graph_id": payload.get("graph_id"),
"success": True,
}
)
except Exception as exc: # pragma: no cover - defensive logging
logger.error(
"Error sending lora code to node %s (graph %s): %s",
parsed_node_id,
graph_identifier,
exc,
)
results.append(
{
"node_id": parsed_node_id,
"graph_id": payload.get("graph_id"),
"success": False,
"error": str(exc),
}
)
return web.json_response({"success": True, "results": results})
except Exception as exc: # pragma: no cover - defensive logging
logger.error("Failed to update lora code (GET): %s", exc, exc_info=True)
return web.json_response({"success": False, "error": str(exc)}, status=500)
class TrainedWordsHandler:
async def get_trained_words(self, request: web.Request) -> web.Response:
@@ -2970,15 +3256,28 @@ class NodeRegistryHandler:
self._node_registry = node_registry
self._prompt_server = prompt_server
self._standalone_mode = standalone_mode
self._refresh_lock = asyncio.Lock()
self._last_slow_path_ts: float = 0.0
async def register_nodes(self, request: web.Request) -> web.Response:
try:
data = await request.json()
nodes = data.get("nodes", [])
client_id = data.get("client_id")
if not isinstance(nodes, list):
return web.json_response(
{"success": False, "error": "nodes must be a list"}, status=400
)
if not isinstance(client_id, str) or not client_id:
return web.json_response(
{
"success": False,
"error": "Missing client_id parameter",
},
status=400,
)
for index, node in enumerate(nodes):
if not isinstance(node, dict):
return web.json_response(
@@ -3005,7 +3304,12 @@ class NodeRegistryHandler:
)
graph_name = node.get("graph_name")
try:
node["node_id"] = int(node_id)
# Handle compound node IDs from expanded group subgraphs,
# e.g. "252:0" → 0 (parent scope is already in graph_id)
if isinstance(node_id, str) and ":" in node_id:
node["node_id"] = int(node_id.rsplit(":", 1)[-1])
else:
node["node_id"] = int(node_id)
except (TypeError, ValueError):
return web.json_response(
{
@@ -3022,7 +3326,7 @@ class NodeRegistryHandler:
else:
node["graph_name"] = str(graph_name)
await self._node_registry.register_nodes(nodes)
await self._node_registry.register_nodes(client_id, nodes)
return web.json_response(
{
"success": True,
@@ -3046,33 +3350,110 @@ class NodeRegistryHandler:
status=503,
)
try:
self._prompt_server.instance.send_sync("lora_registry_refresh", {})
logger.debug("Sent registry refresh request to frontend")
except Exception as exc:
logger.error("Failed to send registry refresh message: %s", exc)
return web.json_response(
{
"success": False,
"error": "Communication Error",
"message": f"Failed to communicate with ComfyUI frontend: {exc}",
},
status=500,
)
current_sids = set(self._prompt_server.instance.sockets.keys())
registry_updated = await self._node_registry.wait_for_update(timeout=1.0)
if not registry_updated:
logger.warning("Registry refresh timeout after 1 second")
# Fast path: if the frontend has already pushed node data (via
# afterConfigureGraph / graphChanged hooks), return it immediately
# without triggering a WebSocket round-trip.
registry_info = await self._node_registry.get_merged_registry(
active_sids=current_sids
)
if registry_info["tab_count"] > 0:
logger.debug(
"[LM:Registry] fast path: %s nodes across %s tabs %s",
registry_info["node_count"],
registry_info["tab_count"],
dict(registry_info.get("tabs", {})),
)
return web.json_response({"success": True, "data": registry_info})
# Slow path: registry is empty — trigger refresh via WebSocket.
# Serialize with an async lock so concurrent callers don't all
# trigger separate WS refresh cycles. The second caller will
# re-check the fast path and (usually) find populated data.
async with self._refresh_lock:
# Re-check after acquiring the lock — another concurrent call
# may have populated the cache while we were waiting.
registry_info = await self._node_registry.get_merged_registry(
active_sids=current_sids
)
if registry_info["tab_count"] > 0:
logger.debug(
"[LM:Registry] fast path after lock wait: %s nodes across %s tabs",
registry_info["node_count"],
registry_info["tab_count"],
)
return web.json_response({"success": True, "data": registry_info})
# Cooldown: if the slow path ran recently (< 2 s) and
# returned empty, skip another WS round-trip.
elapsed = time.monotonic() - self._last_slow_path_ts
if elapsed < 2.0:
logger.debug(
"[LM:Registry] slow path cooldown (%.1fs since last refresh), returning empty",
elapsed,
)
return web.json_response(
{
"success": False,
"error": "Empty Registry",
"message": "No workflow nodes found — ensure ComfyUI is open and the extension is loaded.",
},
status=408,
)
logger.debug(
"[LM:Registry] slow path: cache empty, triggering WS refresh (%s connected tabs: %s)",
len(current_sids), list(current_sids)[:5],
)
active_sids = list(current_sids)
self._node_registry.prepare_for_refresh(active_sids)
try:
self._prompt_server.instance.send_sync("lora_registry_refresh", {})
logger.debug(
"Sent registry refresh request (expecting %s clients)", len(active_sids)
)
except Exception as exc:
logger.error("Failed to send registry refresh message: %s", exc)
return web.json_response(
{
"success": False,
"error": "Communication Error",
"message": f"Failed to communicate with ComfyUI frontend: {exc}",
},
status=500,
)
if not await self._node_registry.wait_for_all(timeout=0.5):
logger.warning(
"Registry refresh timeout after 0.5s (%s/%s clients responded)",
len(active_sids) - self._node_registry.pending_client_count,
len(active_sids),
)
# Re-read current sockets after the wait: a tab may have connected
# while we were waiting, and we don't want to garbage-collect it.
current_sids = set(self._prompt_server.instance.sockets.keys())
registry_info = await self._node_registry.get_merged_registry(
active_sids=current_sids
)
self._last_slow_path_ts = time.monotonic()
if registry_info["node_count"] == 0:
logger.debug(
"[LM:Registry] refresh OK — %s connected tab(s) but 0 compatible nodes found",
registry_info["tab_count"],
)
return web.json_response(
{
"success": False,
"error": "Timeout Error",
"message": "Registry refresh timeout - ComfyUI frontend may not be responsive",
"error": "Empty Registry",
"message": "No workflow nodes found — ensure ComfyUI is open and the extension is loaded.",
},
status=408,
)
registry_info = await self._node_registry.get_registry()
return web.json_response({"success": True, "data": registry_info})
except Exception as exc: # pragma: no cover - defensive logging
logger.error("Failed to get registry: %s", exc, exc_info=True)
@@ -3085,17 +3466,21 @@ class NodeRegistryHandler:
try:
data = await request.json()
widget_name = data.get("widget_name")
action = data.get("action")
value = data.get("value")
mode = data.get("mode", "replace")
node_ids = data.get("node_ids")
if not isinstance(widget_name, str) or not widget_name:
if not action and (not isinstance(widget_name, str) or not widget_name):
return web.json_response(
{"success": False, "error": "Missing widget_name parameter"},
{
"success": False,
"error": "Missing parameter: provide either 'action' or 'widget_name'",
},
status=400,
)
if not isinstance(value, str) or not value:
if value is None or (isinstance(value, str) and not value):
return web.json_response(
{"success": False, "error": "Missing value parameter"}, status=400
)
@@ -3130,12 +3515,15 @@ class NodeRegistryHandler:
except (TypeError, ValueError):
parsed_node_id = node_identifier
payload = {
payload: dict = {
"id": parsed_node_id,
"widget_name": widget_name,
"value": value,
"mode": mode,
}
if action:
payload["action"] = action
if widget_name:
payload["widget_name"] = widget_name
if graph_identifier is not None:
payload["graph_id"] = str(graph_identifier)
@@ -3170,6 +3558,130 @@ class NodeRegistryHandler:
logger.error("Failed to update node widget: %s", exc, exc_info=True)
return web.json_response({"success": False, "error": str(exc)}, status=500)
async def get_update_node_widget(self, request: web.Request) -> web.Response:
"""GET version of update_node_widget — reads parameters from query string.
Query params:
widget_name (optional) the widget name to update (required unless action is set)
action (optional) alternative action, e.g. "inject_text" (required unless widget_name is set)
value (required) the value to set
mode (optional) "replace" (default) or "append"
node_id (repeatable) target node id(s), e.g. node_id=3&node_id=5
node_ids (optional) JSON-encoded array for complex references:
[{"node_id":3,"graph_id":"g1"}, ...]
"""
try:
widget_name = request.query.get("widget_name")
action = request.query.get("action")
value = request.query.get("value")
mode = request.query.get("mode", "replace")
node_ids_raw = request.query.get("node_ids")
node_id_list = request.query.getall("node_id", [])
if not action and (not isinstance(widget_name, str) or not widget_name):
return web.json_response(
{
"success": False,
"error": "Missing parameter: provide either 'action' or 'widget_name'",
},
status=400,
)
if value is None or (isinstance(value, str) and not value):
return web.json_response(
{"success": False, "error": "Missing value parameter"}, status=400
)
node_ids = None
if node_ids_raw:
try:
node_ids = json.loads(node_ids_raw)
except (json.JSONDecodeError, TypeError):
return web.json_response(
{"success": False, "error": "node_ids must be a valid JSON array"},
status=400,
)
if not isinstance(node_ids, list) or not node_ids:
return web.json_response(
{"success": False, "error": "node_ids must be a non-empty JSON array"},
status=400,
)
elif node_id_list:
node_ids = node_id_list
if not isinstance(node_ids, list) or not node_ids:
return web.json_response(
{"success": False, "error": "node_ids must be a non-empty list"},
status=400,
)
results = []
for entry in node_ids:
node_identifier = entry
graph_identifier = None
if isinstance(entry, dict):
node_identifier = entry.get("node_id")
graph_identifier = entry.get("graph_id")
if node_identifier is None:
results.append(
{
"node_id": node_identifier,
"graph_id": graph_identifier,
"success": False,
"error": "Missing node_id parameter",
}
)
continue
try:
parsed_node_id = int(node_identifier)
except (TypeError, ValueError):
parsed_node_id = node_identifier
payload: dict = {
"id": parsed_node_id,
"value": value,
"mode": mode,
}
if action:
payload["action"] = action
if widget_name:
payload["widget_name"] = widget_name
if graph_identifier is not None:
payload["graph_id"] = str(graph_identifier)
try:
self._prompt_server.instance.send_sync("lm_widget_update", payload)
results.append(
{
"node_id": parsed_node_id,
"graph_id": payload.get("graph_id"),
"success": True,
}
)
except Exception as exc: # pragma: no cover - defensive logging
logger.error(
"Error sending widget update to node %s (graph %s): %s",
parsed_node_id,
graph_identifier,
exc,
)
results.append(
{
"node_id": parsed_node_id,
"graph_id": payload.get("graph_id"),
"success": False,
"error": str(exc),
}
)
return web.json_response({"success": True, "results": results})
except Exception as exc: # pragma: no cover - defensive logging
logger.error("Failed to update node widget (GET): %s", exc, exc_info=True)
return web.json_response({"success": False, "error": str(exc)}, status=500)
class MiscHandlerSet:
"""Aggregate handlers into a lookup compatible with the registrar."""
@@ -3194,6 +3706,8 @@ class MiscHandlerSet:
doctor: DoctorHandler,
example_workflows: ExampleWorkflowsHandler,
base_model: BaseModelHandlerSet,
hf_handler: HfHandler | None = None,
agent_handler: AgentHandler | None = None,
) -> None:
self.health = health
self.settings = settings
@@ -3212,6 +3726,8 @@ class MiscHandlerSet:
self.doctor = doctor
self.example_workflows = example_workflows
self.base_model = base_model
self.hf_handler = hf_handler
self.agent_handler = agent_handler
def to_route_mapping(
self,
@@ -3227,13 +3743,17 @@ class MiscHandlerSet:
"get_priority_tags": self.settings.get_priority_tags,
"get_settings_libraries": self.settings.get_libraries,
"activate_library": self.settings.activate_library,
"get_llm_models": self.settings.get_llm_models,
"get_provider_models": self.settings.get_provider_models,
"update_usage_stats": self.usage_stats.update_usage_stats,
"get_usage_stats": self.usage_stats.get_usage_stats,
"update_lora_code": self.lora_code.update_lora_code,
"get_update_lora_code": self.lora_code.get_update_lora_code,
"get_trained_words": self.trained_words.get_trained_words,
"get_model_example_files": self.model_examples.get_model_example_files,
"register_nodes": self.node_registry.register_nodes,
"update_node_widget": self.node_registry.update_node_widget,
"get_update_node_widget": self.node_registry.get_update_node_widget,
"get_registry": self.node_registry.get_registry,
"check_model_exists": self.model_library.check_model_exists,
"check_models_exist": self.model_library.check_models_exist,
@@ -3257,6 +3777,14 @@ class MiscHandlerSet:
"get_supporters": self.supporters.get_supporters,
"get_example_workflows": self.example_workflows.get_example_workflows,
"get_example_workflow": self.example_workflows.get_example_workflow,
# Hugging Face handlers
"get_hf_repo_files": self.hf_handler.get_hf_repo_files,
"download_hf_model": self.hf_handler.download_hf_model,
"set_hf_url": self.hf_handler.set_hf_url,
# Agent skill handlers
"get_agent_skills": self.agent_handler.get_agent_skills,
"execute_agent_skill": self.agent_handler.execute_agent_skill,
"cancel_agent_skill": self.agent_handler.cancel_agent_skill,
# Base model handlers
"get_base_models": self.base_model.get_base_models,
"refresh_base_models": self.base_model.refresh_base_models,
+202 -34
View File
@@ -154,6 +154,14 @@ class ModelPageView:
)
self._template_env._i18n_filter_added = True # type: ignore[attr-defined]
from ...services.llm_service import PROVIDER_PRESETS
# Provider presets are embedded directly (local, no await needed).
# Provider model catalogs are fetched asynchronously by the
# frontend via GET /api/lm/llm/provider-models so page rendering
# never blocks on the remote model catalog (which can take up to
# 30s on cold cache).
template_context = {
"is_initializing": is_initializing,
"settings": self._settings,
@@ -161,6 +169,8 @@ class ModelPageView:
"folders": [],
"t": self._server_i18n.get_translation,
"version": self._get_app_version(),
"provider_presets_json": json.dumps(PROVIDER_PRESETS),
"provider_models_json": "{}",
}
if not is_initializing:
@@ -203,11 +213,17 @@ class ModelListingHandler:
result = await self._service.get_paginated_data(**params)
format_start = time.perf_counter()
formatted_raw = [
await self._service.format_response(entry)
for entry in result["items"]
]
# Filter out None entries returned for corrupted cache rows (issue #730).
# Note: "total" intentionally remains the pre-filter count to reflect
# the true number of models in the cache; corrupted entries are rare
# and adjusting total would cause pagination drift on every page.
formatted_items = [item for item in formatted_raw if item is not None]
formatted_result = {
"items": [
await self._service.format_response(item)
for item in result["items"]
],
"items": formatted_items,
"total": result["total"],
"page": result["page"],
"page_size": result["page_size"],
@@ -233,14 +249,20 @@ class ModelListingHandler:
start_time = time.perf_counter()
try:
params = self._parse_common_params(request)
# group_by_model is meaningless for excluded view; strip it
params.pop("group_by_model", None)
result = await self._service.get_excluded_paginated_data(**params)
format_start = time.perf_counter()
formatted_raw = [
await self._service.format_response(entry)
for entry in result["items"]
]
# Filter out None entries returned for corrupted cache rows (issue #730).
# "total" stays at the pre-filter count; see get_models for rationale.
formatted_items = [item for item in formatted_raw if item is not None]
formatted_result = {
"items": [
await self._service.format_response(item)
for item in result["items"]
],
"items": formatted_items,
"total": result["total"],
"page": result["page"],
"page_size": result["page_size"],
@@ -366,6 +388,21 @@ class ModelListingHandler:
request.query.get("name_pattern_use_regex", "false").lower() == "true"
)
# Group-by-model flag: deduplicate versions sharing the same civitai modelId
group_by_model = (
request.query.get("group_by_model", "false").lower() == "true"
)
# View-local-versions filter: show all local versions of a specific model
# Accepts either a CivitAI modelId (int) or a HF group key like "hf:user/repo"
civitai_model_id = request.query.get("civitai_model_id")
if civitai_model_id is not None:
try:
civitai_model_id = int(civitai_model_id)
except (TypeError, ValueError):
# Keep as string — could be an HF group key (e.g. "hf:user/repo")
pass
return {
"page": page,
"page_size": page_size,
@@ -389,6 +426,8 @@ class ModelListingHandler:
"name_pattern_include": name_pattern_include,
"name_pattern_exclude": name_pattern_exclude,
"name_pattern_use_regex": name_pattern_use_regex,
"group_by_model": group_by_model,
"civitai_model_id": civitai_model_id,
**self._parse_specific_params(request),
}
@@ -500,6 +539,7 @@ class ModelManagementHandler:
# Update model_data with new hash
model_data["sha256"] = sha256
model_data["hash_status"] = "completed"
hash_status = "completed"
else:
return web.json_response(
{"success": False, "error": "No SHA256 hash found"}, status=400
@@ -507,6 +547,32 @@ class ModelManagementHandler:
await MetadataManager.hydrate_model_data(model_data)
# hydrate_model_data replaces model_data with .metadata.json content,
# which may lack sha256. Restore from cache and persist the fix.
if not model_data.get("sha256"):
if sha256:
model_data["sha256"] = sha256
model_data["hash_status"] = model_data.get("hash_status", hash_status)
data_to_save = model_data.copy()
data_to_save.pop("folder", None)
await MetadataManager.save_metadata(file_path, data_to_save)
else:
sha256 = await calculate_sha256(file_path)
if sha256:
model_data["sha256"] = sha256.lower()
model_data["hash_status"] = "completed"
data_to_save = model_data.copy()
data_to_save.pop("folder", None)
await MetadataManager.save_metadata(file_path, data_to_save)
else:
return web.json_response(
{
"success": False,
"error": "Failed to compute SHA256 hash for model",
},
status=500,
)
success, error = await self._metadata_sync.fetch_and_update_model(
sha256=model_data["sha256"],
file_path=file_path,
@@ -516,15 +582,25 @@ class ModelManagementHandler:
if not success:
return web.json_response({"success": False, "error": error})
formatted_metadata = await self._service.format_response(model_data)
return web.json_response({"success": True, "metadata": formatted_metadata})
formatted = await self._service.format_response(model_data)
if formatted is None:
return web.json_response(
{"success": False, "error": "Model entry is corrupted (missing file_path)"},
status=500,
)
return web.json_response({"success": True, "metadata": formatted})
except Exception as exc:
if is_expected_offline_error(str(exc)):
return web.json_response(
{"success": False, "error": OFFLINE_FRIENDLY_MESSAGE},
status=503,
)
self._logger.error("Error fetching from CivitAI: %s", exc, exc_info=True)
self._logger.error(
"Error fetching from CivitAI for %s: %s",
locals().get("file_path", "unknown"),
exc,
exc_info=True,
)
return web.json_response({"success": False, "error": str(exc)}, status=500)
async def relink_civitai(self, request: web.Request) -> web.Response:
@@ -931,6 +1007,8 @@ class ModelQueryHandler:
limit = int(request.query.get("limit", "20"))
if limit < 0:
limit = 20
elif limit > 200:
limit = 20
top_tags = await self._service.get_top_tags(limit)
return web.json_response({"success": True, "tags": top_tags})
except Exception as exc:
@@ -939,6 +1017,22 @@ class ModelQueryHandler:
{"success": False, "error": "Internal server error"}, status=500
)
async def search_tags(self, request: web.Request) -> web.Response:
try:
query = request.query.get("q", "")
limit = int(request.query.get("limit", "20"))
if limit < 0:
limit = 20
elif limit > 200:
limit = 20
tags = await self._service.search_tags(query, limit)
return web.json_response({"success": True, "tags": tags})
except Exception as exc:
self._logger.error("Error searching tags: %s", exc, exc_info=True)
return web.json_response(
{"success": False, "error": "Internal server error"}, status=500
)
async def get_base_models(self, request: web.Request) -> web.Response:
try:
limit = int(request.query.get("limit", "20"))
@@ -1074,10 +1168,12 @@ class ModelQueryHandler:
# Sort: originals first, copies last
sorted_models = self._sort_duplicate_group(filtered)
# Format response
# Format response, filtering out corrupted entries (issue #730)
group = {"hash": sha256, "models": []}
for model in sorted_models:
group["models"].append(await self._service.format_response(model))
formatted = await self._service.format_response(model)
if formatted is not None:
group["models"].append(formatted)
# Only include groups with 2+ models after filtering
if len(group["models"]) > 1:
@@ -1194,9 +1290,9 @@ class ModelQueryHandler:
(m for m in cache.raw_data if m["file_path"] == path), None
)
if model:
group["models"].append(
await self._service.format_response(model)
)
formatted = await self._service.format_response(model)
if formatted is not None:
group["models"].append(formatted)
hash_val = self._service.scanner.get_hash_by_filename(filename)
if hash_val:
main_path = self._service.get_path_by_hash(hash_val)
@@ -1206,9 +1302,9 @@ class ModelQueryHandler:
None,
)
if main_model:
group["models"].insert(
0, await self._service.format_response(main_model)
)
formatted = await self._service.format_response(main_model)
if formatted is not None:
group["models"].insert(0, formatted)
if group["models"]:
result.append(group)
return web.json_response(
@@ -1231,9 +1327,13 @@ class ModelQueryHandler:
text=f"{self._service.model_type.capitalize()} file name is required",
status=400,
)
notes = await self._service.get_model_notes(model_name)
if notes is not None:
return web.json_response({"success": True, "notes": notes})
result = await self._service.get_model_notes(model_name)
if result is not None:
return web.json_response({
"success": True,
"notes": result["notes"],
"file_path": result["file_path"],
})
return web.json_response(
{
"success": False,
@@ -1269,9 +1369,28 @@ class ModelQueryHandler:
}
if include_license_flags:
model_data = await self._service.get_model_info_by_name(model_name)
license_flags = (model_data or {}).get("license_flags")
if license_flags is not None:
response_payload["license_flags"] = int(license_flags)
# Only return license_flags when real CivitAI model license
# data exists. This mirrors ModelModal's guard
# (modelData?.civitai?.model) so the preview tooltip never
# shows misleading license icons for HF or other models
# without actual license metadata.
civitai_data = (model_data or {}).get("civitai") or {}
has_license_data = (
isinstance(civitai_data, dict)
and isinstance(civitai_data.get("model"), dict)
)
if has_license_data:
license_flags = (model_data or {}).get("license_flags")
if license_flags is not None:
response_payload["license_flags"] = int(license_flags)
# Include the user's license icon style preference so the
# ComfyUI tooltip can pick the right set without a separate
# API call.
try:
settings = get_settings_manager()
response_payload["use_new_license_icons"] = settings.get("use_new_license_icons", True)
except Exception:
pass
return web.json_response(response_payload)
return web.json_response(
{
@@ -1720,14 +1839,20 @@ class ModelDownloadHandler:
async def delete_download_history_item(self, request: web.Request) -> web.Response:
try:
item_id = int(request.query.get("id", "0"))
if not item_id:
download_id = request.query.get("download_id")
id_str = request.query.get("id")
item_id = int(id_str) if id_str else None
if not download_id and not item_id:
return web.json_response(
{"success": False, "error": "id is required"}, status=400
{"success": False, "error": "id or download_id is required"},
status=400,
)
service = await DownloadQueueService.get_instance()
deleted = await service.delete_history_item(item_id)
deleted = await service.delete_history_item(
id=item_id, download_id=download_id
)
return web.json_response({"success": deleted})
except Exception as exc:
self._logger.error(
@@ -1737,14 +1862,20 @@ class ModelDownloadHandler:
async def retry_download_from_history(self, request: web.Request) -> web.Response:
try:
item_id = int(request.query.get("id", "0"))
if not item_id:
download_id = request.query.get("download_id")
id_str = request.query.get("id")
item_id = int(id_str) if id_str else None
if not download_id and not item_id:
return web.json_response(
{"success": False, "error": "id is required"}, status=400
{"success": False, "error": "id or download_id is required"},
status=400,
)
service = await DownloadQueueService.get_instance()
item = await service.retry_from_history(item_id)
item = await service.retry_from_history(
item_id=item_id, download_id=download_id
)
if item is None:
return web.json_response(
{"success": False, "error": "History item not found or not retryable"},
@@ -1820,6 +1951,39 @@ class ModelDownloadHandler:
)
return web.json_response({"success": False, "error": str(exc)}, status=500)
async def update_download_queue_status(self, request: web.Request) -> web.Response:
"""Update the status of a queue item (non-terminal transitions).
Supported transitions include ``queued downloading``,
``downloading paused``, ``paused downloading``, etc.
Terminal transitions (``completed``, ``failed``, ``canceled``)
should use ``complete_download_in_queue`` instead.
"""
try:
download_id = request.query.get("download_id")
status = request.query.get("status")
if not download_id or not status:
return web.json_response(
{
"success": False,
"error": "download_id and status are required",
},
status=400,
)
service = await DownloadQueueService.get_instance()
updated = await service.update_status(download_id, status)
if not updated:
return web.json_response(
{"success": False, "error": "Download not found in queue"},
status=404,
)
return web.json_response({"success": True})
except Exception as exc:
self._logger.error(
"Error updating download queue status: %s", exc, exc_info=True
)
return web.json_response({"success": False, "error": str(exc)}, status=500)
class ModelCivitaiHandler:
"""CivitAI integration endpoints."""
@@ -1861,7 +2025,9 @@ class ModelCivitaiHandler:
return web.json_response(result)
except Exception as exc:
self._logger.error(
"Error in fetch_all_civitai for %ss: %s", self._service.model_type, exc
"Error in fetch_all_civitai for %ss: %s",
self._service.model_type, exc,
exc_info=True,
)
return web.Response(text=str(exc), status=500)
@@ -2833,6 +2999,7 @@ class ModelHandlerSet:
"bulk_delete_models": self.management.bulk_delete_models,
"verify_duplicates": self.management.verify_duplicates,
"get_top_tags": self.query.get_top_tags,
"search_tags": self.query.search_tags,
"get_base_models": self.query.get_base_models,
"get_model_types": self.query.get_model_types,
"scan_models": self.query.scan_models,
@@ -2862,6 +3029,7 @@ class ModelHandlerSet:
"retry_all_failed_downloads": self.download.retry_all_failed_downloads,
"complete_download_in_queue": self.download.complete_download_in_queue,
"get_download_stats": self.download.get_download_stats,
"update_download_queue_status": self.download.update_download_queue_status,
"get_civitai_versions": self.civitai.get_civitai_versions,
"get_civitai_model_by_version": self.civitai.get_civitai_model_by_version,
"get_civitai_model_by_hash": self.civitai.get_civitai_model_by_hash,
+31
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
import logging
import mimetypes
import urllib.parse
@@ -53,6 +54,7 @@ class PreviewHandler:
if not resolved.is_file():
logger.debug("Preview file not found at %s", str(resolved))
asyncio.create_task(self._cleanup_stale_preview_url(normalized))
raise web.HTTPNotFound(text="Preview file not found")
# aiohttp's FileResponse handles range requests, content headers, and
@@ -69,6 +71,35 @@ class PreviewHandler:
resp.headers["Cache-Control"] = "public, max-age=86400"
return resp
async def _cleanup_stale_preview_url(self, normalized_preview_path: str) -> None:
"""Fire-and-forget: clear stale preview_url from all model caches.
When a preview file is no longer on disk, remove its reference from
every cached entry so subsequent list API responses return an empty
``preview_url``, letting the frontend show the no-preview placeholder.
"""
try:
from ...services.service_registry import ServiceRegistry
for service_name in ("lora_scanner", "checkpoint_scanner", "embedding_scanner"):
scanner = ServiceRegistry.get_service_sync(service_name)
if scanner is None or not hasattr(scanner, "_cache"):
continue
cache = getattr(scanner, "_cache", None)
if cache is None or not hasattr(cache, "clear_preview_by_path"):
continue
cleared = await cache.clear_preview_by_path(normalized_preview_path)
if cleared and hasattr(scanner, "_persist_current_cache"):
await scanner._persist_current_cache()
logger.info(
"Cleared stale preview_url for %d %s entries (%s)",
cleared,
service_name,
normalized_preview_path,
)
except Exception as exc:
logger.debug("Failed to clean up stale preview_url: %s", exc)
async def _stream_file(
self, request: web.Request, path: Path
) -> web.StreamResponse:
+120 -8
View File
@@ -32,6 +32,7 @@ from ...utils.civitai_utils import (
extract_civitai_image_id_from_cdn_url,
rewrite_preview_url,
)
from ...utils.constants import NSFW_LEVELS
from ...utils.exif_utils import ExifUtils
from ...recipes.merger import GenParamsMerger
from ...recipes.enrichment import RecipeEnricher
@@ -71,6 +72,7 @@ class RecipeHandlerSet:
"save_recipe": self.management.save_recipe,
"delete_recipe": self.management.delete_recipe,
"get_top_tags": self.query.get_top_tags,
"search_tags": self.query.search_tags,
"get_base_models": self.query.get_base_models,
"get_roots": self.query.get_roots,
"get_folders": self.query.get_folders,
@@ -316,12 +318,11 @@ class RecipeQueryHandler:
raise RuntimeError("Recipe scanner unavailable")
limit = int(request.query.get("limit", "20"))
cache = await recipe_scanner.get_cached_data()
tag_counts: Dict[str, int] = {}
for recipe in getattr(cache, "raw_data", []):
for tag in recipe.get("tags", []) or []:
tag_counts[tag] = tag_counts.get(tag, 0) + 1
if limit < 0:
limit = 20
elif limit > 200:
limit = 20
tag_counts = await self._get_recipe_tag_counts(recipe_scanner)
sorted_tags = [
{"tag": tag, "count": count} for tag, count in tag_counts.items()
@@ -332,6 +333,55 @@ class RecipeQueryHandler:
self._logger.error("Error retrieving top tags: %s", exc, exc_info=True)
return web.json_response({"success": False, "error": str(exc)}, status=500)
async def search_tags(self, request: web.Request) -> web.Response:
try:
await self._ensure_dependencies_ready()
recipe_scanner = self._recipe_scanner_getter()
if recipe_scanner is None:
raise RuntimeError("Recipe scanner unavailable")
query = request.query.get("q", "")
limit = int(request.query.get("limit", "20"))
if limit < 0:
limit = 20
elif limit > 200:
limit = 20
tag_counts = await self._get_recipe_tag_counts(recipe_scanner)
normalized_query = (query or "").strip().lower()
if not normalized_query:
sorted_tags = [
{"tag": tag, "count": count} for tag, count in tag_counts.items()
]
sorted_tags.sort(key=lambda entry: entry["count"], reverse=True)
return web.json_response(
{"success": True, "tags": sorted_tags[: (limit if limit > 0 else 20)]}
)
matched = [
{"tag": tag, "count": count}
for tag, count in tag_counts.items()
if normalized_query in tag.lower()
]
matched.sort(key=lambda entry: entry["count"], reverse=True)
if limit == 0:
result = matched
else:
result = matched[:limit]
return web.json_response({"success": True, "tags": result})
except Exception as exc:
self._logger.error("Error searching recipe tags: %s", exc, exc_info=True)
return web.json_response({"success": False, "error": str(exc)}, status=500)
async def _get_recipe_tag_counts(self, recipe_scanner) -> Dict[str, int]:
"""Compute tag->count mapping from cached recipe data."""
cache = await recipe_scanner.get_cached_data()
tag_counts: Dict[str, int] = {}
for recipe in getattr(cache, "raw_data", []):
for tag in recipe.get("tags", []) or []:
tag_counts[tag] = tag_counts.get(tag, 0) + 1
return tag_counts
async def get_base_models(self, request: web.Request) -> web.Response:
try:
await self._ensure_dependencies_ready()
@@ -1120,6 +1170,13 @@ class RecipeManagementHandler:
if parsed_embedded.get("base_model") and not metadata.get("base_model"):
metadata["base_model"] = parsed_embedded["base_model"]
# Extract preview_nsfw_level from the CivitAI API response
# (injected into civitai_meta_raw by _download_remote_media).
if isinstance(civitai_meta_raw, dict):
bl = civitai_meta_raw.get("browsingLevel")
if isinstance(bl, int) and bl > 0:
metadata["preview_nsfw_level"] = bl
civitai_client = self._civitai_client_getter()
await RecipeEnricher.enrich_recipe(
recipe=metadata,
@@ -1515,8 +1572,31 @@ class RecipeManagementHandler:
# CivitAI API returns modelVersionIds at the root level of
# the image response, NOT inside the meta object.
mvids = image_info.get("modelVersionIds")
if mvids and isinstance(civitai_meta_raw, dict):
civitai_meta_raw["modelVersionIds"] = mvids
if mvids:
if isinstance(civitai_meta_raw, dict):
civitai_meta_raw["modelVersionIds"] = mvids
else:
# meta is null but modelVersionIds exists — create a
# minimal dict so downstream parsers can discover
# LoRAs and checkpoints from the API response.
civitai_meta_raw = {"modelVersionIds": mvids}
# Inject browsingLevel (canonical integer) so the recipe's
# preview_nsfw_level can be set, enabling proper NSFW blur
# of the preview image. Fall back to nsfwLevel (string)
# when browsingLevel is absent.
if isinstance(civitai_meta_raw, dict):
browsing_level = image_info.get("browsingLevel")
nsfw_level_str = image_info.get("nsfwLevel")
if isinstance(browsing_level, int) and browsing_level > 0:
civitai_meta_raw["browsingLevel"] = browsing_level
elif (
isinstance(nsfw_level_str, str)
and nsfw_level_str in NSFW_LEVELS
):
civitai_meta_raw["browsingLevel"] = NSFW_LEVELS[
nsfw_level_str
]
original_url = (
image_info.get("url") if civitai_image_id and image_info else None
@@ -1796,6 +1876,13 @@ class RecipeManagementHandler:
"source_path": image_url,
}
# Extract preview_nsfw_level from the CivitAI API response
# (injected into civitai_meta_raw by _download_remote_media).
if isinstance(civitai_meta_raw, dict):
bl = civitai_meta_raw.get("browsingLevel")
if isinstance(bl, int) and bl > 0:
metadata["preview_nsfw_level"] = bl
if civitai_parsed:
civitai_loras = civitai_parsed.get("loras", [])
if civitai_loras and not metadata.get("loras"):
@@ -2180,6 +2267,31 @@ class RecipeManagementHandler:
"Failed to download image for recipe: %s", exc
)
# Fallback: try to locate a custom image on disk using model_hash + image id
if image_bytes is None:
image_id = image_data.get("id") or ""
if image_id and model_hash:
from ...utils.example_images_paths import get_model_folder
model_folder = get_model_folder(model_hash)
if model_folder and os.path.exists(model_folder):
for fname in os.listdir(model_folder):
if f"custom_{image_id}" in fname:
ext = os.path.splitext(fname)[1].lower()
if ext not in (".jpg", ".jpeg", ".png", ".webp", ".gif"):
continue
fpath = os.path.join(model_folder, fname)
if os.path.isfile(fpath):
try:
with open(fpath, "rb") as f:
image_bytes = f.read()
extension = ext
except Exception as exc:
self._logger.warning(
"Failed to read custom image file %s: %s",
fpath, exc,
)
break
prompt = (
(parsed.get("gen_params") or {}).get("prompt") or ""
)
+24
View File
@@ -22,6 +22,8 @@ class RouteDefinition:
MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition("GET", "/api/lm/settings", "get_settings"),
RouteDefinition("POST", "/api/lm/settings", "update_settings"),
RouteDefinition("GET", "/api/lm/llm/models", "get_llm_models"),
RouteDefinition("GET", "/api/lm/llm/provider-models", "get_provider_models"),
RouteDefinition("GET", "/api/lm/doctor/diagnostics", "get_doctor_diagnostics"),
RouteDefinition("POST", "/api/lm/doctor/repair-cache", "repair_doctor_cache"),
RouteDefinition("POST", "/api/lm/doctor/resolve-filename-conflicts", "resolve_doctor_filename_conflicts"),
@@ -37,10 +39,12 @@ MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition("POST", "/api/lm/update-usage-stats", "update_usage_stats"),
RouteDefinition("GET", "/api/lm/get-usage-stats", "get_usage_stats"),
RouteDefinition("POST", "/api/lm/update-lora-code", "update_lora_code"),
RouteDefinition("GET", "/api/lm/update-lora-code", "get_update_lora_code"),
RouteDefinition("GET", "/api/lm/trained-words", "get_trained_words"),
RouteDefinition("GET", "/api/lm/model-example-files", "get_model_example_files"),
RouteDefinition("POST", "/api/lm/register-nodes", "register_nodes"),
RouteDefinition("POST", "/api/lm/update-node-widget", "update_node_widget"),
RouteDefinition("GET", "/api/lm/update-node-widget", "get_update_node_widget"),
RouteDefinition("GET", "/api/lm/get-registry", "get_registry"),
RouteDefinition("GET", "/api/lm/check-model-exists", "check_model_exists"),
RouteDefinition("GET", "/api/lm/check-models-exist", "check_models_exist"),
@@ -94,6 +98,26 @@ MISC_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition(
"GET", "/api/lm/delete-model-version", "delete_model_version"
),
# Hugging Face model endpoints
RouteDefinition(
"GET", "/api/lm/hf-repo-files", "get_hf_repo_files"
),
RouteDefinition(
"POST", "/api/lm/download-hf-model", "download_hf_model"
),
RouteDefinition(
"POST", "/api/lm/set-hf-url", "set_hf_url"
),
# Agent skill endpoints
RouteDefinition(
"GET", "/api/lm/agent/skills", "get_agent_skills"
),
RouteDefinition(
"POST", "/api/lm/agent/execute/{skill_name}", "execute_agent_skill"
),
RouteDefinition(
"POST", "/api/lm/agent/cancel", "cancel_agent_skill"
),
)
+6
View File
@@ -39,6 +39,8 @@ from .handlers.misc_handlers import (
build_service_registry_adapter,
)
from .handlers.base_model_handlers import BaseModelHandlerSet
from .handlers.hf_handlers import HfHandler
from .handlers.agent_handlers import AgentHandler
from .misc_route_registrar import MiscRouteRegistrar
logger = logging.getLogger(__name__)
@@ -136,6 +138,8 @@ class MiscRoutes:
doctor = DoctorHandler(settings_service=self._settings)
example_workflows = ExampleWorkflowsHandler()
base_model = BaseModelHandlerSet()
hf_handler = HfHandler()
agent_handler = AgentHandler()
return self._handler_set_factory(
health=health,
@@ -155,6 +159,8 @@ class MiscRoutes:
doctor=doctor,
example_workflows=example_workflows,
base_model=base_model,
hf_handler=hf_handler,
agent_handler=agent_handler,
)
+4
View File
@@ -46,6 +46,7 @@ COMMON_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
"GET", "/api/lm/{prefix}/auto-organize-progress", "get_auto_organize_progress"
),
RouteDefinition("GET", "/api/lm/{prefix}/top-tags", "get_top_tags"),
RouteDefinition("GET", "/api/lm/{prefix}/search-tags", "search_tags"),
RouteDefinition("GET", "/api/lm/{prefix}/base-models", "get_base_models"),
RouteDefinition("GET", "/api/lm/{prefix}/model-types", "get_model_types"),
RouteDefinition("GET", "/api/lm/{prefix}/scan", "scan_models"),
@@ -138,6 +139,9 @@ COMMON_ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition(
"GET", "/api/lm/downloads/queue/complete", "complete_download_in_queue"
),
RouteDefinition(
"GET", "/api/lm/downloads/queue/status", "update_download_queue_status"
),
RouteDefinition("POST", "/api/lm/{prefix}/cancel-task", "cancel_task"),
RouteDefinition("GET", "/{prefix}", "handle_models_page"),
)
+1
View File
@@ -29,6 +29,7 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition("POST", "/api/lm/recipes/save", "save_recipe"),
RouteDefinition("DELETE", "/api/lm/recipe/{recipe_id}", "delete_recipe"),
RouteDefinition("GET", "/api/lm/recipes/top-tags", "get_top_tags"),
RouteDefinition("GET", "/api/lm/recipes/search-tags", "search_tags"),
RouteDefinition("GET", "/api/lm/recipes/base-models", "get_base_models"),
RouteDefinition("GET", "/api/lm/recipes/roots", "get_roots"),
RouteDefinition("GET", "/api/lm/recipes/folders", "get_folders"),
+45 -16
View File
@@ -11,6 +11,8 @@ from ..config import config
from ..services.settings_manager import get_settings_manager
from ..services.server_i18n import server_i18n
from ..services.service_registry import ServiceRegistry
from ..services.model_query import normalize_sub_type, resolve_sub_type
from ..utils.constants import VALID_LORA_SUB_TYPES, VALID_CHECKPOINT_SUB_TYPES
from ..utils.usage_stats import UsageStats
logger = logging.getLogger(__name__)
@@ -140,6 +142,21 @@ class StatsRoutes:
# Get usage statistics
usage_data = await self.usage_stats.get_stats()
# CivitAI model type distribution across all model types
# Use the same logic as the filter panel: normalize_sub_type(resolve_sub_type(entry))
# with sub-type validation per model type
model_types_counter: Counter[str] = Counter()
for entry in lora_cache.raw_data:
ntype = normalize_sub_type(resolve_sub_type(entry))
if ntype and ntype in VALID_LORA_SUB_TYPES:
model_types_counter[ntype] += 1
for entry in checkpoint_cache.raw_data:
ntype = normalize_sub_type(resolve_sub_type(entry))
if ntype and ntype in VALID_CHECKPOINT_SUB_TYPES:
model_types_counter[ntype] += 1
# Embeddings: always count as "embedding" regardless of CivitAI sub-type
model_types_counter['embedding'] = len(embedding_cache.raw_data)
return web.json_response({
'success': True,
'data': {
@@ -154,7 +171,8 @@ class StatsRoutes:
'total_generations': usage_data.get('total_executions', 0),
'unused_loras': self._count_unused_models(lora_cache.raw_data, usage_data.get('loras', {})),
'unused_checkpoints': self._count_unused_models(checkpoint_cache.raw_data, usage_data.get('checkpoints', {})),
'unused_embeddings': self._count_unused_models(embedding_cache.raw_data, usage_data.get('embeddings', {}))
'unused_embeddings': self._count_unused_models(embedding_cache.raw_data, usage_data.get('embeddings', {})),
'model_types_distribution': dict(model_types_counter.most_common())
}
})
@@ -459,9 +477,12 @@ class StatsRoutes:
if unused_lora_percent > 50:
insights.append({
'type': 'warning',
'title': 'High Number of Unused LoRAs',
'description': f'{unused_lora_percent:.1f}% of your LoRAs ({unused_loras}/{total_loras}) have never been used.',
'suggestion': 'Consider organizing or archiving unused models to free up storage space.'
'key': 'insights.unusedLoras.high',
'params': {
'percent': f'{unused_lora_percent:.1f}',
'count': str(unused_loras),
'total': str(total_loras)
}
})
if total_checkpoints > 0:
@@ -469,9 +490,12 @@ class StatsRoutes:
if unused_checkpoint_percent > 30:
insights.append({
'type': 'warning',
'title': 'Unused Checkpoints Detected',
'description': f'{unused_checkpoint_percent:.1f}% of your checkpoints ({unused_checkpoints}/{total_checkpoints}) have never been used.',
'suggestion': 'Review and consider removing checkpoints you no longer need.'
'key': 'insights.unusedCheckpoints.detected',
'params': {
'percent': f'{unused_checkpoint_percent:.1f}',
'count': str(unused_checkpoints),
'total': str(total_checkpoints)
}
})
if total_embeddings > 0:
@@ -479,9 +503,12 @@ class StatsRoutes:
if unused_embedding_percent > 50:
insights.append({
'type': 'warning',
'title': 'High Number of Unused Embeddings',
'description': f'{unused_embedding_percent:.1f}% of your embeddings ({unused_embeddings}/{total_embeddings}) have never been used.',
'suggestion': 'Consider organizing or archiving unused embeddings to optimize your collection.'
'key': 'insights.unusedEmbeddings.high',
'params': {
'percent': f'{unused_embedding_percent:.1f}',
'count': str(unused_embeddings),
'total': str(total_embeddings)
}
})
# Storage insights
@@ -492,18 +519,20 @@ class StatsRoutes:
if total_size > 100 * 1024 * 1024 * 1024: # 100GB
insights.append({
'type': 'info',
'title': 'Large Collection Detected',
'description': f'Your model collection is using {self._format_size(total_size)} of storage.',
'suggestion': 'Consider using external storage or cloud solutions for better organization.'
'key': 'insights.collection.large',
'params': {
'size': self._format_size(total_size)
}
})
# Recent activity insight
if usage_data.get('total_executions', 0) > 100:
insights.append({
'type': 'success',
'title': 'Active User',
'description': f'You\'ve completed {usage_data["total_executions"]} generations so far!',
'suggestion': 'Keep exploring and creating amazing content with your models.'
'key': 'insights.activity.active',
'params': {
'count': str(usage_data['total_executions'])
}
})
return web.json_response({
+342 -49
View File
@@ -16,6 +16,105 @@ logger = logging.getLogger(__name__)
NETWORK_EXCEPTIONS = (ClientError, OSError, asyncio.TimeoutError)
# User-managed directories that live inside the plugin folder (portable
# mode) and must survive a Git-based update. ``git clean -fd`` would
# otherwise delete them because they are untracked and, in released tags,
# not listed in ``.gitignore``. ``-e`` excludes a path from cleaning
# regardless of whether it is ignored.
_PRESERVE_DIRS = ('settings.json', 'civitai', 'wildcards', 'backups', 'stats', 'logs', 'cache', 'model_cache')
def _clean_excludes() -> List[str]:
"""Build the ``-e`` arguments for ``git clean`` from :data:`_PRESERVE_DIRS`."""
excludes: List[str] = []
for name in _PRESERVE_DIRS:
excludes.append('-e')
excludes.append(name)
# For directories, also exclude nested matches explicitly
# (``-e dir`` alone matches the dir entry; ``-e dir/**`` guards
# contents under all git versions as defense-in-depth).
excludes.append('-e')
excludes.append(f'{name}/**')
return excludes
def _stage_preserved_items(plugin_root: str) -> tuple[str, list[str]]:
"""Move preserved user-data items to a temp directory outside *plugin_root*.
This ensures that ``git reset --hard``, ``git clean -fd``, and ZIP-based
replacement cannot touch these files even when ``-e`` exclusion patterns
are mishandled (e.g. on Windows where forward-slash patterns may not
match backslash-prefixed paths in some Git builds, or where file locks
prevent deletion/recreation).
Returns:
``(backup_root, staged_names)``: the temp directory path and the
list of item names that were successfully moved.
"""
backup_root = tempfile.mkdtemp(prefix='lora_manager_update_')
staged: list[str] = []
for name in _PRESERVE_DIRS:
src = os.path.join(plugin_root, name)
if not os.path.lexists(src):
continue
dst = os.path.join(backup_root, name)
try:
shutil.move(src, dst)
staged.append(name)
logger.debug("Staged '%s' for update safety", name)
except OSError:
# ``shutil.move`` may fail on Windows if a file handle inside
# the directory is still open (e.g. a SQLite WAL file). Fall
# back to copy-then-remove.
logger.debug("Move failed for '%s', falling back to copy", name)
try:
if os.path.isdir(src) and not os.path.islink(src):
shutil.copytree(src, dst, symlinks=True)
shutil.rmtree(src, ignore_errors=True)
else:
shutil.copy2(src, dst)
os.remove(src)
staged.append(name)
logger.info("Copied (then removed) '%s' for update safety", name)
except Exception as exc:
logger.warning(
"Could not stage '%s': %s (will rely on git -e / skip lists)", name, exc
)
return backup_root, staged
def _restore_preserved_items(plugin_root: str, backup_root: str, staged: list[str]) -> None:
"""Move staged items back from *backup_root* into *plugin_root*.
Any leftover placeholder at the destination (created by git checkout or
ZIP extraction) is removed before the move.
"""
for name in staged:
src = os.path.join(backup_root, name)
dst = os.path.join(plugin_root, name)
try:
if os.path.lexists(dst):
if os.path.isdir(dst) and not os.path.islink(dst):
shutil.rmtree(dst, ignore_errors=True)
else:
os.remove(dst)
shutil.move(src, dst)
logger.debug("Restored '%s' after update", name)
except OSError:
logger.debug("Move failed restoring '%s', falling back to copy", name)
try:
if os.path.isdir(src) and not os.path.islink(src):
shutil.copytree(src, dst, symlinks=True, dirs_exist_ok=True)
shutil.rmtree(src, ignore_errors=True)
else:
shutil.copy2(src, dst)
os.remove(src)
logger.info("Copied '%s' back after update", name)
except Exception as exc:
logger.error("Failed to restore '%s': %s", name, exc)
shutil.rmtree(backup_root, ignore_errors=True)
class UpdateRoutes:
"""Routes for handling plugin update checks"""
@@ -26,6 +125,7 @@ class UpdateRoutes:
app.router.add_get('/api/lm/check-updates', UpdateRoutes.check_updates)
app.router.add_get('/api/lm/version-info', UpdateRoutes.get_version_info)
app.router.add_post('/api/lm/perform-update', UpdateRoutes.perform_update)
app.router.add_post('/api/lm/switch-channel', UpdateRoutes.switch_channel)
@staticmethod
async def check_updates(request):
@@ -44,10 +144,17 @@ class UpdateRoutes:
# Fetch remote version from GitHub
if nightly:
remote_version, changelog = await UpdateRoutes._get_nightly_version()
releases = None
local_hash = git_info.get('short_hash', '')
nightly_version, releases_result = await asyncio.gather(
UpdateRoutes._get_nightly_version(local_hash),
UpdateRoutes._get_remote_version()
)
remote_version, _, behind_by, commit_date = nightly_version
_, changelog, releases = releases_result
else:
remote_version, changelog, releases = await UpdateRoutes._get_remote_version()
behind_by = 0
commit_date = ''
# Compare versions
if nightly:
@@ -60,6 +167,10 @@ class UpdateRoutes:
remote_version.replace('v', '')
)
current_dir = os.path.dirname(os.path.abspath(__file__))
plugin_root = os.path.dirname(os.path.dirname(current_dir))
has_git = os.path.exists(os.path.join(plugin_root, '.git'))
response_data = {
'success': True,
'current_version': local_version,
@@ -67,13 +178,13 @@ class UpdateRoutes:
'update_available': update_available,
'changelog': changelog,
'git_info': git_info,
'nightly': nightly
'nightly': nightly,
'has_git': has_git,
'releases': releases,
'behind_by': behind_by,
'commit_date': commit_date
}
# Include releases list for stable mode
if releases is not None:
response_data['releases'] = releases
return web.json_response(response_data)
except NETWORK_EXCEPTIONS as e:
@@ -105,9 +216,14 @@ class UpdateRoutes:
# Format: version-short_hash
version_string = f"{local_version}-{short_hash}"
current_dir = os.path.dirname(os.path.abspath(__file__))
plugin_root = os.path.dirname(os.path.dirname(current_dir))
has_git = os.path.exists(os.path.join(plugin_root, '.git'))
return web.json_response({
'success': True,
'version': version_string
'version': version_string,
'has_git': has_git
})
except Exception as e:
@@ -135,20 +251,22 @@ class UpdateRoutes:
if os.path.exists(settings_path):
with open(settings_path, 'r', encoding='utf-8') as f:
settings_backup = f.read()
logger.info("Backed up settings.json")
logger.debug("Backed up settings.json (%d bytes)", len(settings_backup))
git_folder = os.path.join(plugin_root, '.git')
if os.path.exists(git_folder):
# Git update
success, new_version = await UpdateRoutes._perform_git_update(plugin_root, nightly)
else:
# Fallback: Download ZIP and replace files
success, new_version = await UpdateRoutes._download_and_replace_zip(plugin_root)
staged_backup_dir, staged_items = _stage_preserved_items(plugin_root)
try:
git_folder = os.path.join(plugin_root, '.git')
if os.path.exists(git_folder):
success, new_version = await UpdateRoutes._perform_git_update(plugin_root, nightly)
else:
success, new_version = await UpdateRoutes._download_and_replace_zip(plugin_root)
finally:
_restore_preserved_items(plugin_root, staged_backup_dir, staged_items)
if settings_backup and success:
with open(settings_path, 'w', encoding='utf-8') as f:
f.write(settings_backup)
logger.info("Restored settings.json")
logger.debug("Restored settings.json content (%d bytes)", len(settings_backup))
if success:
return web.json_response({
@@ -169,6 +287,164 @@ class UpdateRoutes:
'error': str(e)
})
@staticmethod
async def switch_channel(request):
"""
Switch between release and nightly update channels.
ZIP/CNR install Nightly: git init + checkout main (one-way upgrade)
Git install Release: git checkout latest tag (.git preserved)
ZIP/CNR install Release: ZIP download (no .git, stays in ZIP mode)
Git install Nightly: git checkout main + pull
"""
try:
body = await request.json() if request.has_body else {}
channel = body.get('channel', '')
if channel not in ('release', 'nightly'):
return web.json_response({
'success': False,
'error': f'Invalid channel: {channel}. Must be "release" or "nightly".'
})
current_dir = os.path.dirname(os.path.abspath(__file__))
plugin_root = os.path.dirname(os.path.dirname(current_dir))
settings_path = ensure_settings_file(logger)
settings_backup = None
if os.path.exists(settings_path):
with open(settings_path, 'r', encoding='utf-8') as f:
settings_backup = f.read()
logger.debug("Backed up settings.json before channel switch (%d bytes)", len(settings_backup))
staged_backup_dir, staged_items = _stage_preserved_items(plugin_root)
try:
git_folder = os.path.join(plugin_root, '.git')
if channel == 'nightly':
git_backup = None
if os.path.exists(git_folder):
git_backup = UpdateRoutes._backup_git(git_folder, 'nightly')
success = False
new_version = ''
try:
if os.path.exists(git_folder):
success, new_version = await UpdateRoutes._perform_git_update(
plugin_root, nightly=True
)
else:
success, new_version = UpdateRoutes._init_git_repo(plugin_root)
finally:
UpdateRoutes._restore_git(git_backup, git_folder, success, 'nightly')
else:
success = False
new_version = ''
if os.path.exists(git_folder):
success, new_version = await UpdateRoutes._perform_git_update(
plugin_root, nightly=False
)
else:
tracking_file = os.path.join(plugin_root, '.tracking')
if os.path.exists(tracking_file):
os.remove(tracking_file)
success, new_version = await UpdateRoutes._download_and_replace_zip(plugin_root)
finally:
_restore_preserved_items(plugin_root, staged_backup_dir, staged_items)
if settings_backup and success:
with open(settings_path, 'w', encoding='utf-8') as f:
f.write(settings_backup)
logger.debug("Restored settings.json content after channel switch (%d bytes)", len(settings_backup))
if success:
return web.json_response({
'success': True,
'channel': channel,
'new_version': new_version,
'message': f'Switched to {channel} channel'
})
else:
return web.json_response({
'success': False,
'error': f'Failed to switch to {channel} channel'
})
except Exception as e:
logger.error("Failed to switch channel: %s", e, exc_info=True)
return web.json_response({
'success': False,
'error': str(e)
})
@staticmethod
def _init_git_repo(plugin_root: str) -> tuple[bool, str]:
"""
Initialize a Git repository in a ZIP-installed plugin folder.
Clones the remote history and checks out main branch.
"""
try:
import git
except ImportError:
logger.error(
"GitPython is not available: cannot initialize git repo. "
"Install git or set $GIT_PYTHON_GIT_EXECUTABLE to the git binary path."
)
return False, ""
clean_excludes = _clean_excludes()
try:
repo = git.Repo.init(plugin_root)
origin = repo.create_remote(
'origin',
'https://github.com/willmiao/ComfyUI-Lora-Manager.git'
)
origin.fetch()
repo.create_head('main', origin.refs.main)
repo.git.checkout('main', '--force')
repo.git.reset('--hard')
repo.git.clean('-fd', *clean_excludes)
tracking_file = os.path.join(plugin_root, '.tracking')
if os.path.exists(tracking_file):
os.remove(tracking_file)
logger.info("Removed .tracking file (now in git mode)")
new_version = f"main-{repo.head.commit.hexsha[:7]}"
logger.info("Initialized git repo on main branch: %s", new_version)
return True, new_version
except Exception as e:
logger.error("Failed to initialize git repo: %s", e, exc_info=True)
return False, ""
@staticmethod
def _backup_git(git_folder, label):
try:
backup_dir = tempfile.mkdtemp()
backup = os.path.join(backup_dir, '.git')
shutil.copytree(git_folder, backup)
logger.info("Backed up .git before switching to %s", label)
return backup
except Exception as e:
logger.error("Failed to backup .git before %s switch: %s", label, e)
return None
@staticmethod
def _restore_git(git_backup, git_folder, success, label):
if git_backup and not success:
try:
if os.path.exists(git_folder):
shutil.rmtree(git_folder)
shutil.copytree(git_backup, git_folder)
logger.info("Restored .git after failed %s switch", label)
except Exception as e:
logger.error("Failed to restore .git after %s switch: %s", label, e)
if git_backup:
shutil.rmtree(os.path.dirname(git_backup), ignore_errors=True)
@staticmethod
async def _download_and_replace_zip(plugin_root: str) -> tuple[bool, str]:
"""
@@ -223,8 +499,7 @@ class UpdateRoutes:
except Exception:
logger.debug("Could not close downloaded-version history database", exc_info=True)
# Skip settings.json, civitai, model cache and runtime cache folders
UpdateRoutes._clean_plugin_folder(plugin_root, skip_files=['settings.json', 'civitai', 'model_cache', 'cache', 'wildcards', 'backups', 'stats'])
UpdateRoutes._clean_plugin_folder(plugin_root, skip_files=list(_PRESERVE_DIRS))
# Extract ZIP to temp dir
with tempfile.TemporaryDirectory() as tmp_dir:
@@ -234,7 +509,7 @@ class UpdateRoutes:
extracted_root = next(os.scandir(tmp_dir)).path
# Copy files, skipping user data that should be preserved
skip_items = {'settings.json', 'civitai', 'wildcards', 'backups', 'stats'}
skip_items = set(_PRESERVE_DIRS)
for item in os.listdir(extracted_root):
if item in skip_items:
continue
@@ -251,7 +526,7 @@ class UpdateRoutes:
# for ComfyUI Manager to work properly
tracking_info_file = os.path.join(plugin_root, '.tracking')
tracking_files = []
skip_tracked = {'civitai', 'wildcards', 'backups', 'stats'}
skip_tracked = set(_PRESERVE_DIRS) - {'settings.json'}
for root, dirs, files in os.walk(extracted_root):
# Skip user data directories and their contents
rel_root = os.path.relpath(root, extracted_root)
@@ -274,7 +549,8 @@ class UpdateRoutes:
except Exception as e:
logger.error(f"ZIP update failed: {e}", exc_info=True)
return False, ""
@staticmethod
def _clean_plugin_folder(plugin_root, skip_files=None):
skip_files = skip_files or []
for item in os.listdir(plugin_root):
@@ -287,41 +563,54 @@ class UpdateRoutes:
os.remove(path)
@staticmethod
async def _get_nightly_version() -> tuple[str, List[str]]:
"""
Fetch latest commit from main branch
"""
async def _get_nightly_version(local_hash: str = "") -> tuple[str, List[str], int, str]:
repo_owner = "willmiao"
repo_name = "ComfyUI-Lora-Manager"
# Use GitHub API to fetch the latest commit from main branch
github_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/commits/main"
try:
downloader = await get_downloader()
success, data = await downloader.make_request('GET', github_url, custom_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:
logger.warning(f"Failed to fetch GitHub commit: {data}")
return "main", []
commit_sha = data.get('sha', '')[:7] # Short hash
logger.warning("Failed to fetch GitHub commit: %s", data)
return "main", [], 0, ""
commit_sha = data.get('sha', '')[:7]
commit_message = data.get('commit', {}).get('message', '')
# Format as "main-{short_hash}"
commit_date = data.get('commit', {}).get('committer', {}).get('date', '')[:10]
version = f"main-{commit_sha}"
# Use commit message as changelog
changelog = [commit_message] if commit_message else []
return version, changelog
behind_by = 0
if local_hash and local_hash not in ('unknown', 'stable'):
compare_url = (
f"https://api.github.com/repos/{repo_owner}/{repo_name}"
f"/compare/{local_hash}...main"
)
c_ok, c_data = await downloader.make_request(
'GET', compare_url,
custom_headers={'Accept': 'application/vnd.github+json'}
)
if c_ok:
if c_data.get('status') in ('ahead', 'diverged'):
behind_by = c_data.get('ahead_by', 0)
else:
behind_by = c_data.get('behind_by', 0)
return version, changelog, behind_by, commit_date
except NETWORK_EXCEPTIONS as e:
logger.warning("Unable to reach GitHub for nightly version: %s", e)
return "main", []
return "main", [], 0, ""
except Exception as e:
logger.error(f"Error fetching nightly version: {e}", exc_info=True)
return "main", []
logger.error("Error fetching nightly version: %s", e, exc_info=True)
return "main", [], 0, ""
@staticmethod
def _compare_nightly_versions(local_git_info: Dict[str, str], remote_version: str) -> bool:
@@ -365,6 +654,8 @@ class UpdateRoutes:
)
return False, ""
clean_excludes = _clean_excludes()
try:
# Open the Git repository
repo = git.Repo(plugin_root)
@@ -376,8 +667,9 @@ class UpdateRoutes:
if nightly:
# Reset to discard any local changes
repo.git.reset('--hard')
# Clean untracked files
repo.git.clean('-fd')
# Clean untracked files, but preserve user-managed directories
# (wildcards, backups, stats, civitai, caches, settings.json).
repo.git.clean('-fd', *clean_excludes)
# Switch to main branch and pull latest
main_branch = 'main'
@@ -394,8 +686,9 @@ class UpdateRoutes:
else:
# Reset to discard any local changes
repo.git.reset('--hard')
# Clean untracked files
repo.git.clean('-fd')
# Clean untracked files, but preserve user-managed directories
# (wildcards, backups, stats, civitai, caches, settings.json).
repo.git.clean('-fd', *clean_excludes)
# Get latest release tag
tags = sorted(repo.tags, key=lambda t: t.commit.committed_datetime, reverse=True)
+27
View File
@@ -0,0 +1,27 @@
"""LLM-powered metadata enrichment pipeline infrastructure.
This package provides the orchestration layer for LLM-powered features.
Skills define *what* to do (prompt template). The :class:`AgentService`
handles *how* (LLM calls, context gathering, validation, progress).
NOTE: The current implementation is a code-driven pipeline, not a true
agent loop. Future agent orchestration (LLM-driven tool selection) will
live alongside this package with its own namespace.
"""
from __future__ import annotations
from .skill_definition import SkillDefinition, SkillPermissions
from .skill_registry import SkillRegistry
from .agent_service import AgentService, AgentProgressReporter, SkillResult
from .post_processor import PostProcessor
__all__ = [
"AgentProgressReporter",
"AgentService",
"PostProcessor",
"SkillDefinition",
"SkillPermissions",
"SkillRegistry",
"SkillResult",
]
+489
View File
@@ -0,0 +1,489 @@
"""Pipeline orchestration service.
The :class:`AgentService` coordinates LLM-powered pipeline execution:
1. Look up the pipeline definition in :class:`SkillRegistry`
2. Validate input against its ``input_schema``
3. Prepare context via :mod:`~py.metadata_ops` (read metadata, list base models, fetch HF README)
4. If ``llm_required``: call :class:`LLMService` with the rendered prompt
5. Post-process via :class:`PostProcessor` (delegates I/O to :mod:`~py.metadata_ops`)
6. Broadcast progress and completion via :class:`WebSocketManager`
Pipeline definitions (*skills*) describe *what* to do (prompt template).
The AgentService handles *how* (LLM calls, context gathering, validation,
progress).
"""
from __future__ import annotations
import asyncio
import json
import logging
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
import aiohttp
import os
from ...config import config
from ..llm_service import LLMService
from ..websocket_manager import ws_manager
from .post_processor import PostProcessor
from .skill_registry import SkillRegistry
from .skills.enrich_hf_metadata.readme_processor import (
clean_readme_for_llm,
extract_relevant_section,
)
logger = logging.getLogger(__name__)
class AgentProgressReporter:
"""Protocol-compatible progress reporter backed by WebSocket broadcast."""
async def on_progress(self, payload: Dict[str, Any]) -> None:
await ws_manager.broadcast(payload)
@dataclass
class SkillResult:
"""Outcome of a skill execution."""
success: bool
updated_models: List[Dict[str, Any]] = field(default_factory=list)
errors: List[str] = field(default_factory=list)
summary: str = ""
def _validate_schema(data: Any, schema: Dict[str, Any], path: str = "") -> List[str]:
"""Minimal JSON schema validator.
Supports a subset of JSON Schema: ``type``, ``properties``, ``required``,
``items``, ``enum``. Returns a list of error messages (empty = valid).
"""
errors: List[str] = []
if not schema:
return errors
expected_type = schema.get("type")
if expected_type:
type_map = {
"string": str,
"number": (int, float),
"integer": int,
"boolean": bool,
"array": list,
"object": dict,
"null": type(None),
}
expected_py = type_map.get(expected_type)
if expected_py is not None and not isinstance(data, expected_py):
errors.append(f"{path or 'root'}: expected {expected_type}, got {type(data).__name__}")
return errors
if expected_type == "object" and isinstance(data, dict):
properties = schema.get("properties", {})
required = schema.get("required", [])
for req_key in required:
if req_key not in data:
errors.append(f"{path or 'root'}: missing required property '{req_key}'")
for key, value in data.items():
if key in properties:
errors.extend(_validate_schema(value, properties[key], f"{path}.{key}"))
if expected_type == "array" and isinstance(data, list):
items_schema = schema.get("items")
if items_schema:
for i, item in enumerate(data):
errors.extend(_validate_schema(item, items_schema, f"{path}[{i}]"))
if "enum" in schema and data not in schema["enum"]:
errors.append(f"{path or 'root'}: value '{data}' not in enum {schema['enum']}")
return errors
# ------------------------------------------------------------------
# Prompt template rendering
# ------------------------------------------------------------------
def _render_prompt(template: str, variables: Dict[str, Any]) -> str:
"""Render a prompt template with ``{{variable}}`` placeholders.
Uses simple regex substitution no Jinja2 dependency needed.
"""
def replace(match: re.Match) -> str:
key = match.group(1).strip()
value = variables.get(key, "")
if isinstance(value, (dict, list)):
return json.dumps(value, ensure_ascii=False, indent=2)
return str(value)
return re.sub(r"\{\{(\w+)\}\}", replace, template)
class AgentService:
"""Orchestrate agent skill execution.
Usage::
service = await AgentService.get_instance()
result = await service.execute_skill(
skill_name="enrich_hf_metadata",
input_data={"model_paths": ["/path/to/model.safetensors"]},
progress_callback=AgentProgressReporter(),
)
"""
_instance: Optional["AgentService"] = None
_lock: asyncio.Lock = asyncio.Lock()
def __init__(
self,
*,
skill_registry: Optional[SkillRegistry] = None,
llm_service: Optional[LLMService] = None,
) -> None:
self._registry = skill_registry
self._llm_service = llm_service
@classmethod
async def get_instance(cls) -> "AgentService":
"""Return the lazily-initialised global ``AgentService``."""
if cls._instance is None:
async with cls._lock:
if cls._instance is None:
cls._instance = cls(
skill_registry=await SkillRegistry.get_instance(),
llm_service=await LLMService.get_instance(),
)
return cls._instance
@classmethod
def reset_instance(cls) -> None:
"""Reset the cached singleton — primarily for tests."""
cls._instance = None
async def _ensure_registry(self) -> SkillRegistry:
if self._registry is None:
self._registry = await SkillRegistry.get_instance()
return self._registry
async def _ensure_llm(self) -> LLMService:
if self._llm_service is None:
self._llm_service = await LLMService.get_instance()
return self._llm_service
async def list_skills(self) -> List[Dict[str, Any]]:
"""Return a JSON-serialisable list of available skills."""
registry = await self._ensure_registry()
return [
{
"name": s.name,
"title": s.title,
"description": s.description,
"llm_required": s.llm_required,
"model_type_filter": s.model_type_filter,
}
for s in registry.list_skills()
]
async def execute_skill(
self,
*,
skill_name: str,
input_data: Dict[str, Any],
progress_callback: Optional[AgentProgressReporter] = None,
) -> SkillResult:
"""Execute a pipeline (skill) on the given models.
Args:
skill_name: Name of the pipeline to execute
input_data: Input validated against the pipeline's ``input_schema``
progress_callback: Optional WebSocket progress reporter
Returns:
:class:`SkillResult` with success status and updated model info
"""
registry = await self._ensure_registry()
skill = registry.get_skill(skill_name)
if skill is None:
return SkillResult(
success=False,
errors=[f"Skill not found: {skill_name}"],
summary=f"Skill '{skill_name}' does not exist",
)
input_errors = _validate_schema(input_data, skill.input_schema)
if input_errors:
return SkillResult(
success=False,
errors=input_errors,
summary=f"Invalid input: {'; '.join(input_errors)}",
)
model_paths = input_data.get("model_paths", [])
if not model_paths:
return SkillResult(
success=False,
errors=["No model_paths provided"],
summary="No models to process",
)
total = len(model_paths)
processed = 0
success_count = 0
skipped_count = 0
updated_models: List[Dict[str, Any]] = []
errors: List[str] = []
post_processor = PostProcessor()
await self._emit_progress(
progress_callback, skill_name, status="started",
total=total, processed=0, success=0,
)
llm = await self._ensure_llm()
llm_configured = llm.is_configured() if skill.llm_required else True
for model_path in model_paths:
model_filename = os.path.basename(model_path)
logger.info(
"[%s] [%d/%d] %s",
skill_name, processed + 1, total, model_filename,
)
updated_data: Dict[str, Any] = {}
skip_model = False
try:
from ...metadata_ops import read_metadata
metadata = await read_metadata(model_path)
# Fast-fail: enrich_hf_metadata requires hf_url to have HF README context
if skill_name == "enrich_hf_metadata" and not metadata.get("hf_url", ""):
logger.info(
"[%s] SKIP %s — no hf_url in metadata",
skill_name, model_filename,
)
skipped_count += 1
skip_model = True
if not skip_model:
prompt_vars: Dict[str, Any] = {"model_path": model_path}
if skill.llm_required and llm_configured:
prompt_vars = await self._build_prompt_context(
skill_name, model_path, metadata, registry, llm,
)
llm_response: Optional[Dict[str, Any]] = None
if skill.llm_required and llm_configured:
prompt_template = registry.load_prompt(skill_name)
rendered = _render_prompt(prompt_template, prompt_vars)
llm_response = await llm.chat_completion_json(
system_prompt=prompt_vars.get(
"system_prompt",
"You are a helpful assistant that extracts structured metadata.",
),
user_prompt=rendered,
)
if llm_response:
logger.info(
"[%s] [%d/%d] %s → base_model=%s confidence=%s",
skill_name, processed + 1, total, model_filename,
(llm_response.get("base_model") or "?")[:50],
llm_response.get("confidence", "?"),
)
model_result = await post_processor.process(
skill_name=skill_name,
model_path=model_path,
llm_output=llm_response or {},
metadata=metadata,
readme_content=prompt_vars.get("readme_content_full", ""),
)
if model_result.get("success", True):
success_count += 1
uf = model_result.get("updated_fields", [])
if uf:
updated_models.append({"path": model_path, "updated_fields": uf})
updated_data = model_result.get("updates", {})
if "preview_url" in updated_data and updated_data["preview_url"]:
updated_data["preview_url"] = config.get_preview_static_url(
updated_data["preview_url"]
)
else:
errors.extend(
model_result.get("errors", [model_result.get("error", "Unknown error")])
)
except Exception as exc:
logger.error("Skill %s failed for %s: %s", skill_name, model_path, exc)
errors.append(f"{model_path}: {exc}")
processed += 1
await self._emit_progress(
progress_callback, skill_name, status="processing",
total=total, processed=processed, success=success_count,
skipped=skipped_count,
current_path=model_path,
updated_data=updated_data,
)
result = SkillResult(
success=success_count > 0,
updated_models=updated_models,
errors=errors,
summary=f"Processed {processed}/{total} models, {success_count} succeeded, {skipped_count} skipped",
)
await self._emit_progress(
progress_callback, skill_name, status="completed",
total=total, processed=processed, success=success_count,
skipped=skipped_count,
updated_models=updated_models, errors=errors, summary=result.summary,
)
return result
# ------------------------------------------------------------------
# Base model grouping (keeps the prompt compact)
# ------------------------------------------------------------------
@staticmethod
def _format_base_models(models: List[str]) -> str:
"""Format the base model list as a flat, one-per-line list.
Attempts to group by family consistently degraded LLM extraction
accuracy the LLM finds individual model names harder to spot
in comma-separated groups than in a simple ``- Name`` list.
"""
return "\n".join(f"- {m}" for m in models)
async def _build_prompt_context(
self,
skill_name: str,
model_path: str,
metadata: Dict[str, Any],
registry: SkillRegistry,
llm: Any,
) -> Dict[str, Any]:
"""Gather variables for the skill's prompt template.
Reads metadata, fetches the HF README (if applicable), lists available
base models, loads user priority tags, and returns a dict that maps to
``{{variable}}`` placeholders in ``prompt.md``.
"""
from ...metadata_ops import identify_model_type, list_base_models
from ..settings_manager import SettingsManager
context: Dict[str, Any] = {
"model_path": model_path,
"model_basename": "",
"hf_url": "",
"repo": "",
"readme_content": "",
"readme_content_full": "",
"current_metadata": {},
"base_models": [],
"priority_tags": "",
}
# Extract model basename (filename without extension) for the LLM
# to use when locating the matching section in collection repos.
raw_basename = os.path.splitext(os.path.basename(model_path))[0]
context["model_basename"] = raw_basename or ""
context["current_metadata"] = {
"file_name": metadata.get("file_name", ""),
"base_model": metadata.get("base_model", ""),
"tags": metadata.get("tags", []),
"modelDescription": metadata.get("modelDescription", ""),
"trainedWords": metadata.get("trainedWords", []),
"sha256": (metadata.get("sha256") or "")[:16] + "..." if metadata.get("sha256") else "",
"size": metadata.get("size", 0),
}
hf_url = metadata.get("hf_url", "")
context["hf_url"] = hf_url
repo = self._extract_repo_from_url(hf_url) if hf_url else ""
context["repo"] = repo or ""
if repo:
readme = await self._fetch_readme(repo)
# Trim README to the section relevant to this model file
# (collection repos often have multiple models in one README).
if readme and raw_basename:
trimmed = extract_relevant_section(readme, raw_basename)
cleaned = clean_readme_for_llm(trimmed) if trimmed else ""
else:
cleaned = clean_readme_for_llm(readme) if readme else ""
context["readme_content"] = cleaned if cleaned else "(README not available)"
context["readme_content_full"] = readme or ""
try:
raw_models = await list_base_models()
context["base_models"] = self._format_base_models(raw_models)
except Exception as exc:
logger.debug("Failed to list base models: %s", exc)
context["base_models"] = "</not available>"
# Determine model type and load the corresponding priority_tags
try:
model_type = await identify_model_type(model_path)
context["model_type"] = model_type
settings = SettingsManager()
priority_config = settings.get_priority_tag_config()
context["priority_tags"] = priority_config.get(model_type, "")
except Exception as exc:
logger.debug("Failed to load priority tags: %s", exc)
context["model_type"] = "lora"
context["priority_tags"] = ""
return context
@staticmethod
def _extract_repo_from_url(hf_url: str) -> Optional[str]:
"""Extract ``user/repo`` from a HuggingFace URL."""
if not hf_url:
return None
m = re.match(r"https?://huggingface\.co/([^/]+/[^/]+)", hf_url)
return m.group(1) if m else None
@staticmethod
async def _fetch_readme(repo: str) -> str:
"""Fetch README.md from HuggingFace (tries ``main``, then ``master``)."""
async with aiohttp.ClientSession(
headers={"User-Agent": "ComfyUI-LoRA-Manager/1.0"},
timeout=aiohttp.ClientTimeout(total=30),
) as session:
for branch in ("main", "master"):
url = f"https://huggingface.co/{repo}/raw/{branch}/README.md"
try:
async with session.get(url) as resp:
if resp.status == 200:
return await resp.text()
except Exception as exc:
logger.debug("Failed to fetch README from %s: %s", url, exc)
return ""
async def _emit_progress(
self,
callback: Optional[AgentProgressReporter],
skill_name: str,
*,
status: str,
**extra: Any,
) -> None:
"""Send a progress update via WebSocket (if callback is set)."""
payload: Dict[str, Any] = {"type": "agent_progress", "skill": skill_name, "status": status}
payload.update(extra)
if callback is not None:
await callback.on_progress(payload)
+336
View File
@@ -0,0 +1,336 @@
"""Post-processing engine for skill pipeline outputs.
The :class:`PostProcessor` takes the LLM's structured JSON output and applies
it to a model's on-disk metadata via the :mod:`~py.metadata_ops` functions.
It handles all the skill-specific business logic conditions, transformations,
and orchestration of multiple side-effects (write metadata, download preview,
refresh cache). All actual I/O is delegated to :mod:`~py.metadata_ops`.
"""
from __future__ import annotations
import json
import logging
import os
import re
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
logger = logging.getLogger(__name__)
class PostProcessor:
"""Deterministic post-processor for skill pipeline outputs.
Usage (called by :class:`~py.services.agent.agent_service.AgentService`)::
processor = PostProcessor()
result = await processor.process(
skill_name="enrich_hf_metadata",
model_path="/path/to/model.safetensors",
llm_output={...},
metadata={...}, # from metadata_ops.read_metadata()
)
"""
async def process(
self,
*,
skill_name: str,
model_path: str,
llm_output: Dict[str, Any],
metadata: Dict[str, Any],
readme_content: str = "",
) -> Dict[str, Any]:
"""Route *llm_output* to the correct skill post-processor.
*readme_content* is optional raw markdown content (e.g. HF README)
that is converted to HTML and stored as ``modelDescription`` for
the description tab.
Returns a dict with keys ``success`` (bool), ``updated_fields`` (list),
``preview_downloaded`` (bool), and ``errors`` (list).
"""
if skill_name == "enrich_hf_metadata":
return await self._process_enrich_hf_metadata(
model_path, llm_output, metadata, readme_content,
)
return {
"success": False,
"updated_fields": [],
"errors": [f"No post-processor registered for skill: {skill_name}"],
}
# ------------------------------------------------------------------
# enrich_hf_metadata
# ------------------------------------------------------------------
async def _process_enrich_hf_metadata(
self,
model_path: str,
llm_output: Dict[str, Any],
metadata: Dict[str, Any],
readme_content: str = "",
) -> Dict[str, Any]:
from ...metadata_ops import (
apply_metadata_updates,
download_preview,
refresh_cache,
)
from .skills.enrich_hf_metadata.readme_processor import (
convert_readme_to_html,
extract_gallery_images,
extract_gallery_table_images,
extract_relevant_section,
extract_simple_markdown_images,
extract_html_img_tags,
extract_repo_from_hf_url,
)
updated_fields: List[str] = []
preview_downloaded = False
# -- Determine whether this is an HF-sourced model -----------------
is_hf_model = not metadata.get("from_civitai", True)
# -- Collect updates -----------------------------------------------
updates: Dict[str, Any] = {}
# base_model
new_base = (llm_output.get("base_model") or "").strip()
current_base = metadata.get("base_model", "") or ""
if new_base and self._should_overwrite(current_base, is_hf_model):
updates["base_model"] = new_base
# trigger words → civitai.trainedWords
new_triggers = llm_output.get("trigger_words", [])
trigger_words_empty = True
if isinstance(new_triggers, list):
cleaned = [t.strip() for t in new_triggers if t.strip()]
cleaned = [t for t in cleaned if t.lower() not in ("none", "null", "n/a")]
trigger_words_empty = not cleaned
current_civitai = metadata.get("civitai") or {}
current_triggers = current_civitai.get("trainedWords") or []
if self._should_overwrite_list(current_triggers, is_hf_model):
trig_civitai = dict(current_civitai)
if "civitai" in updates and isinstance(updates["civitai"], dict):
trig_civitai.update(updates["civitai"])
trig_civitai["trainedWords"] = cleaned
updates["civitai"] = trig_civitai
# modelDescription — from raw README content (converted to HTML)
if readme_content and is_hf_model:
converted = convert_readme_to_html(readme_content)
if converted:
updates["modelDescription"] = converted
# short_description → civitai.description (for "About this version")
short_desc = (llm_output.get("short_description") or "").strip()
if short_desc and is_hf_model:
current_civitai = metadata.get("civitai") or {}
desc_civitai = dict(current_civitai)
if "civitai" in updates and isinstance(updates["civitai"], dict):
desc_civitai.update(updates["civitai"])
desc_civitai["description"] = short_desc
updates["civitai"] = desc_civitai
# gallery images → civitai.images (from YAML frontmatter widget entries
# and Sample Gallery markdown tables in the README body)
gallery_images: List[Dict[str, Any]] = []
if readme_content and is_hf_model:
hf_url = metadata.get("hf_url", "") or ""
repo = extract_repo_from_hf_url(hf_url)
if repo:
rec_w = llm_output.get("recommended_width") or 0
rec_h = llm_output.get("recommended_height") or 0
# 1. Widget images (YAML frontmatter)
gallery = extract_gallery_images(
readme_content, repo,
default_width=rec_w, default_height=rec_h,
)
# 2. Sample Gallery table images (markdown body), deduplicated
existing_urls = {img["url"] for img in gallery if img.get("url")}
table_images = extract_gallery_table_images(
readme_content, repo,
existing_urls=existing_urls,
default_width=rec_w, default_height=rec_h,
)
existing_urls.update(img["url"] for img in table_images if img.get("url"))
# 3. Simple markdown images `![alt](url)` in the body
simple_images = extract_simple_markdown_images(
readme_content, repo,
existing_urls=existing_urls,
default_width=rec_w, default_height=rec_h,
)
existing_urls.update(img["url"] for img in simple_images if img.get("url"))
# 4. HTML `<img>` tags (used by many collection repos)
html_images = extract_html_img_tags(
readme_content, repo,
existing_urls=existing_urls,
default_width=rec_w, default_height=rec_h,
)
all_images = gallery + table_images + simple_images + html_images
if all_images:
gallery_images = all_images
current_civitai = metadata.get("civitai") or {}
gallery_civitai = dict(current_civitai)
if "civitai" in updates and isinstance(updates["civitai"], dict):
gallery_civitai.update(updates["civitai"])
gallery_civitai["images"] = all_images
updates["civitai"] = gallery_civitai
# tags
new_tags = llm_output.get("tags", [])
if isinstance(new_tags, list) and new_tags:
existing_tags = metadata.get("tags") or []
merged = self._merge_tags(existing_tags, new_tags)
if len(merged) > len(existing_tags) or is_hf_model:
updates["tags"] = merged
# metadata_source & llm_enriched_at (always set)
updates["metadata_source"] = "agent:enrich_hf_metadata"
updates["llm_enriched_at"] = datetime.now(timezone.utc).isoformat()
# Store LLM confidence in metadata so it's accessible for evaluation
raw_confidence = (llm_output.get("confidence") or "").strip()
if raw_confidence:
updates["_llm_confidence"] = raw_confidence
# Fallback: extract instance_prompt from YAML frontmatter when the LLM
# returned empty trigger words but the README has instance_prompt.
if trigger_words_empty:
instance_prompt = _extract_yaml_instance_prompt(readme_content)
if instance_prompt:
current_civitai = metadata.get("civitai") or {}
trig_civitai = dict(current_civitai)
if "civitai" in updates and isinstance(updates["civitai"], dict):
trig_civitai.update(updates["civitai"])
trig_civitai["trainedWords"] = [instance_prompt]
updates["civitai"] = trig_civitai
preview_remote_url = (llm_output.get("preview_url") or "").strip()
# Fallback: if the LLM couldn't find a preview image in the cleaned
# README, find the first gallery image from the *model-specific
# section* of the README (not the repo-wide first image, which
# belongs to a different model in collection repos).
if not preview_remote_url and readme_content and is_hf_model:
model_basename = os.path.splitext(os.path.basename(model_path))[0]
relevant_section = extract_relevant_section(
readme_content, model_basename,
)
if relevant_section and relevant_section != readme_content:
for img in gallery_images:
img_url = img.get("url", "")
if img_url and img_url in relevant_section:
preview_remote_url = img_url
break
# Last resort: use the first gallery image from the full README.
if not preview_remote_url and gallery_images:
preview_remote_url = gallery_images[0].get("url", "")
current_preview = metadata.get("preview_url") or ""
if preview_remote_url and not (current_preview and os.path.exists(current_preview)):
local_path = await download_preview(model_path, preview_remote_url)
if local_path:
preview_downloaded = True
updates["preview_url"] = local_path
# notes — plain-text summary of usage info from the LLM
new_notes = (llm_output.get("notes") or "").strip()
if new_notes:
updates["notes"] = new_notes
# usage_tips — JSON string (e.g. {"strength_min":0.85,"strength_max":1.4})
raw_tips = (llm_output.get("usage_tips") or "").strip()
if raw_tips and raw_tips != "{}":
try:
json.loads(raw_tips)
updates["usage_tips"] = raw_tips
except (json.JSONDecodeError, TypeError):
logger.warning(
"LLM returned invalid usage_tips JSON: %s", raw_tips[:200]
)
if updates:
updated_fields = await apply_metadata_updates(model_path, updates)
# -- Refresh scanner cache ------------------------------------------
if updated_fields or preview_downloaded:
await refresh_cache(model_path)
return {
"success": True,
"updated_fields": updated_fields,
"preview_downloaded": preview_downloaded,
"updates": updates,
"errors": [],
}
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
@staticmethod
def _should_overwrite(current_value: str, is_hf_model: bool) -> bool:
"""Return ``True`` when a scalar field should be overwritten."""
return is_hf_model or not current_value or current_value.lower() in (
"", "unknown",
)
@staticmethod
def _should_overwrite_list(current_list: List[str], is_hf_model: bool) -> bool:
"""Return ``True`` when a list field should be overwritten."""
return is_hf_model or not current_list
@staticmethod
def _merge_tags(existing: List[str], new: List[str]) -> List[str]:
"""Merge *new* tags into *existing*, all lowercased.
This matches the behaviour of :class:`TagUpdateService` which
normalises every tag to lowercase for case-insensitive dedup.
"""
merged: List[str] = []
seen: set = set()
for tag in list(existing) + list(new):
t = tag.strip().lower()
if t and t not in seen:
merged.append(t)
seen.add(t)
return merged
# ------------------------------------------------------------------
# Module-level helpers
# ------------------------------------------------------------------
def _extract_yaml_instance_prompt(readme_content: str) -> str:
"""Extract ``instance_prompt`` from the YAML frontmatter of a HF README.
Returns the prompt text, or empty string if not found. Handles
``null`` / ``~`` YAML null values by returning empty string.
"""
if not readme_content or not readme_content.startswith("---"):
return ""
# Find end of frontmatter
end = readme_content.find("---", 3)
if end == -1:
return ""
frontmatter = readme_content[3:end]
for line in frontmatter.split("\n"):
line = line.strip()
m = re.match(r"^instance_prompt:\s*(.*)", line)
if m:
val = m.group(1).strip().strip('"').strip("'")
if val.lower() in ("null", "~", "none", ""):
return ""
return val
return ""
+45
View File
@@ -0,0 +1,45 @@
"""Skill definition data structures.
Each skill is described by a :class:`SkillDefinition` that declares its
input/output schemas, whether it needs an LLM call, and what permissions
its post-processor has.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
@dataclass(frozen=True)
class SkillPermissions:
"""Declarative permission scope for a skill's post-processor.
These are auditable constraints the :class:`AgentService` checks them
before invoking the handler. They are defense-in-depth, not a sandbox.
"""
write_metadata: bool = True
write_previews: bool = True
network_domains: Tuple[str, ...] = ()
@dataclass(frozen=True)
class SkillDefinition:
"""Immutable description of an agent skill."""
name: str
title: str
description: str
llm_required: bool
input_schema: Dict[str, Any] = field(default_factory=dict)
output_schema: Dict[str, Any] = field(default_factory=dict)
model_type_filter: Optional[List[str]] = None
permissions: SkillPermissions = field(default_factory=SkillPermissions)
def applies_to_model_type(self, model_type: str) -> bool:
"""Return ``True`` if this skill can run on the given model type."""
if self.model_type_filter is None:
return True
return model_type in self.model_type_filter
+210
View File
@@ -0,0 +1,210 @@
"""Discovery and loading of prompt-based skills.
Skills live in ``py/services/agent/skills/<name>/`` directories. Each
directory must contain a ``prompt.md`` file with YAML frontmatter::
---
name: my_skill
title: "My Skill"
description: "What this skill does"
llm_required: true
---
Prompt template with ``{{variable}}`` placeholders.
Legacy ``SKILL.md`` files are also supported for backward compatibility.
The registry scans the skills directory on first access and caches results.
"""
from __future__ import annotations
import asyncio
import logging
import re
from pathlib import Path
from typing import Any, Dict, List, Optional
import yaml
from .skill_definition import SkillDefinition, SkillPermissions
logger = logging.getLogger(__name__)
# Directory where built-in skills are stored
_SKILLS_DIR = Path(__file__).parent / "skills"
#: Preferred file names for prompt definition files (tried in order).
#: ``prompt.md`` is the current convention; ``SKILL.md`` is the legacy name
#: kept for backward compatibility.
_PROMPT_FILE_NAMES: tuple[str, ...] = ("prompt.md", "SKILL.md")
# ---------------------------------------------------------------------------
# Frontmatter parser
# ---------------------------------------------------------------------------
_FRONTMATTER_RE = re.compile(
r"^---\s*\n(.*?\n)---\s*\n?(.*)", re.DOTALL
)
def _parse_skill_file(path: Path) -> tuple[dict, str]:
"""Read a prompt definition file (``prompt.md`` or legacy ``SKILL.md``) and
return (frontmatter_dict, body_text).
Raises ``ValueError`` if the file lacks valid YAML frontmatter.
"""
text = path.read_text(encoding="utf-8")
m = _FRONTMATTER_RE.match(text)
if not m:
raise ValueError(f"Missing or invalid YAML frontmatter in {path}")
frontmatter = yaml.safe_load(m.group(1))
if not isinstance(frontmatter, dict):
raise ValueError(f"Frontmatter in {path} is not a mapping")
body = m.group(2).strip()
return frontmatter, body
class SkillRegistry:
"""Discover and load agent skills from the filesystem."""
_instance: Optional["SkillRegistry"] = None
_lock: asyncio.Lock = asyncio.Lock()
def __init__(self, skills_dir: Path = _SKILLS_DIR) -> None:
self._skills_dir = skills_dir
self._skills: Dict[str, SkillDefinition] = {}
self._loaded: bool = False
# ------------------------------------------------------------------
# Singleton access
# ------------------------------------------------------------------
@classmethod
async def get_instance(cls) -> "SkillRegistry":
"""Return the lazily-initialised global ``SkillRegistry``."""
if cls._instance is None:
async with cls._lock:
if cls._instance is None:
registry = cls()
registry._discover()
cls._instance = registry
return cls._instance
@classmethod
def reset_instance(cls) -> None:
"""Reset the cached singleton — primarily for tests."""
cls._instance = None
# ------------------------------------------------------------------
# Discovery
# ------------------------------------------------------------------
@staticmethod
def _find_prompt_file(skill_dir: Path) -> Path | None:
"""Return the first prompt definition file that exists in *skill_dir*.
Tries ``_PROMPT_FILE_NAMES`` in order so that new conventions
(``prompt.md``) take precedence while legacy ``SKILL.md`` files
still load without changes.
"""
for name in _PROMPT_FILE_NAMES:
candidate = skill_dir / name
if candidate.exists():
return candidate
return None
def _discover(self) -> None:
"""Scan the skills directory and load all valid skill definitions."""
self._skills.clear()
if not self._skills_dir.is_dir():
logger.warning("Skills directory does not exist: %s", self._skills_dir)
self._loaded = True
return
for entry in sorted(self._skills_dir.iterdir()):
if not entry.is_dir():
continue
prompt_file = self._find_prompt_file(entry)
if prompt_file is None:
continue
try:
definition = self._load_skill_definition(prompt_file)
if definition is not None:
self._skills[definition.name] = definition
logger.debug("Loaded skill: %s", definition.name)
except Exception as exc:
logger.warning("Failed to load skill from %s: %s", prompt_file, exc)
self._loaded = True
logger.info("Discovered %d prompt-based skills", len(self._skills))
def _load_skill_definition(self, path: Path) -> Optional[SkillDefinition]:
"""Parse a prompt definition file's frontmatter into a
:class:`SkillDefinition`."""
try:
data, _body = _parse_skill_file(path)
except (ValueError, yaml.YAMLError) as exc:
logger.warning("Failed to parse prompt file %s: %s", path, exc)
return None
if "name" not in data:
logger.warning("Prompt file %s missing required 'name' field", path)
return None
perm_data = data.get("permissions", {})
permissions = SkillPermissions(
write_metadata=perm_data.get("write_metadata", True),
write_previews=perm_data.get("write_previews", True),
network_domains=tuple(perm_data.get("network_domains", [])),
)
return SkillDefinition(
name=data["name"],
title=data.get("title", data["name"]),
description=data.get("description", ""),
llm_required=data.get("llm_required", False),
input_schema=data.get("input_schema", {}),
output_schema=data.get("output_schema", {}),
model_type_filter=data.get("model_type_filter"),
permissions=permissions,
)
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def list_skills(self) -> List[SkillDefinition]:
"""Return all discovered skill definitions."""
if not self._loaded:
self._discover()
return list(self._skills.values())
def get_skill(self, name: str) -> Optional[SkillDefinition]:
"""Return the skill definition for ``name``, or ``None`` if not found."""
if not self._loaded:
self._discover()
return self._skills.get(name)
def load_prompt(self, name: str) -> str:
"""Load and return the prompt template body for the named skill."""
skill_dir = self._skills_dir / name
skill_path = self._find_prompt_file(skill_dir)
if skill_path is None:
raise FileNotFoundError(
f"Prompt file not found for skill '{name}' in {skill_dir} "
f"(tried {list(_PROMPT_FILE_NAMES)})"
)
try:
_frontmatter, body = _parse_skill_file(skill_path)
return body
except (ValueError, yaml.YAMLError) as exc:
raise ValueError(f"Failed to parse prompt from {skill_path}: {exc}") from exc
@@ -0,0 +1,165 @@
---
name: enrich_hf_metadata
title: "Enrich Metadata from HuggingFace"
description: >
Parse the HuggingFace model card via LLM to extract description, trigger
words, base model, tags, and preview image URL.
llm_required: true
---
You are an expert assistant for AI image generation models. Your task is to extract structured metadata from a HuggingFace model card (README.md).
## Model Information
- **Repository**: {{hf_url}}
- **Model file path**: {{model_path}}
- **Model filename**: {{model_basename}}
- **Repository ID**: {{repo}}
## Current Metadata (may be incomplete)
```json
{{current_metadata}}
```
## User Priority Tags Reference
The user has configured the following list of **meaningful tag categories** for this model type (`{{model_type}}`):
```
{{priority_tags}}
```
These are the subjects, styles, and concepts the user considers useful for categorization. Use this list as a **reference** when evaluating tags (see the **tags** section below).
## Available Base Models
The following base models are currently valid in this system. Use the EXACT
name listed — do not invent aliases or modify variant suffixes.
{{base_models}}
## HuggingFace README Content
```
{{readme_content}}
```
## Extraction Instructions
Extract the following information from the README content above:
### base_model
The base model this model was trained on. Use EXACTLY one of the names from the **Available Base Models** list above. Do not invent new names or use aliases.
Check the YAML frontmatter for ``base_model:`` first. If the frontmatter has no ``base_model:``, look at the **model filename** (``{{model_basename}}``), YAML ``tags:``, README title and first paragraph for clues — the base model family is often embedded in the name
### trigger_words
The trigger words or activation prompts needed to use this LoRA. Look for:
- `instance_prompt:` in the YAML frontmatter
- Phrases like "trigger word:", "trigger:", "use this prompt:", "activation prompt:"
- In collection repos: the trigger section **specific to this model file** (look near matching download links or anchor IDs)
- Example prompts at the start (usually the first word or phrase before any description)
Return as an array of strings. If none found, return an empty array `[]`. **Never** return `["None"]` or any placeholder value — a truly empty list means no trigger words exist.
### short_description
A concise 1-2 sentence summary of what this model does. Extract from the "Model description" section or the first paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`, not the repo as a whole. Return empty string if the README is too minimal.
### tags
3-8 relevant tags for categorizing this model. **Quality over quantity.**
Sources to consider:
- The YAML frontmatter `tags:` list (filter out technical ones — see below)
- The subject, style, character, or concept the model represents
- The model filename itself may give clues (e.g. "pokemon", "anime", "pixelart")
**Critical filtering rules — apply them strictly:**
1. **Exclude technical/generic tags.** Reject any tag that describes the model's **training methodology, framework, architecture, or modality** rather than its content. Examples to exclude: `text-to-image`, `diffusers`, `lora`, `dreambooth`, `diffusers-training`, `flux`, `sdxl`, `checkpoint`, `pytorch`, `safetensors`, `fine-tuning`, `stable-diffusion`, and any variant of these.
2. **Cross-reference against the priority_tags reference.** Only include a tag if it meaningfully describes what the model actually creates (subject, style, character type) and is semantically close to one of the priority_tags. If none of the README's tags match meaningful categories, prefer returning a smaller set or an empty array over including low-value tags.
3. **All lowercase, no spaces, no hyphens** (use single words like `"photorealistic"`, `"anime"`, `"character"`).
Return empty array if no meaningful content tags remain after filtering.
### recommended_width, recommended_height
The recommended image generation resolution for this model, in pixels. Look for sections like "Best Dimensions", "Recommended size", "Suggested resolution", or similar phrasing in the README. Prefer the explicitly marked "Best" or default resolution. If the table/list has multiple entries (e.g. "768 x 1024 (Best)" and "1024 x 1024 (Default)"), use the one marked "Best". Return integers. If no resolution can be determined, return 0 for both.
### preview_url
The URL of the most suitable preview image from the README. Look for:
- Image tags near the section matching the model filename (`{{model_basename}}`)
- The YAML frontmatter `widget:` section (which often has `output.url` fields)
- In collection repos: the sample images listed **under the section** for this specific model version
- Generic `![alt](url)` in the body
Choose the first image that appears to be a generation example (not a logo or diagram). Construct the absolute URL as `https://huggingface.co/{{repo}}/resolve/main/{filename}`. If no suitable image is found, return an empty string.
### notes
A plain-text summary of the model card's key practical usage information. Combine trigger words, style modifiers, recommended parameters (steps, CFG, resolution, sampler), and any setup tips into a readable paragraph. For collection repos, focus on the **specific model version** matching `{{model_basename}}`. Return empty string if the README has no useful usage info.
### usage_tips
A JSON string with structured usage recommendations. Extract from the README any explicit ranges or recommended values (e.g. "Set LoRA strength: **0.85 - 1.4**", "CLIP strength: 0.5"). Possible fields (include only those you can determine):
```json
{
"strength_min": 0.85,
"strength_max": 1.4,
"strength_range": "0.85-1.4",
"strength": 0.6,
"clip_strength": 0.5,
"clip_skip": 2
}
```
Return the JSON string (e.g. `'{"strength_min":0.85,"strength_max":1.4}'`). Return `"{}"` if nothing useful is found.
### confidence
Your confidence level in the extracted data:
- "high" — most fields were explicitly stated in the README
- "medium" — some fields were inferred from context
- "low" — most fields are guesses based on limited information
## Important: Handling Collection Repos (multiple model files)
Many HuggingFace repos contain **multiple model files** in a single repository
(e.g. a "LoRA collection" with different styles/characters in separate files).
The model file currently being enriched is: **`{{model_basename}}`**
To find the correct section in the README:
1. **Search for download links** containing the filename — the surrounding paragraph is your section.
2. **Search for anchor IDs** (`<a id="...">`) or section headings whose text matches words from the filename.
3. **Search for HTML headings** (`<h1>`, `<h2>`, `<span>`) containing parts of the filename.
4. If no match is found, use the full README as usual — the model may be the only one in the repo.
When a matching section IS found, prefer metadata from that section.
When no section matches (e.g. single-model repos or repos without per-file sections),
extract metadata from the full README normally. Do not return empty data just
because the filename doesn't appear in the README.
## Output Format
Return ONLY a JSON object with exactly these fields (no markdown fences, no extra text):
```json
{
"model_path": "{{model_path}}",
"base_model": "<canonical name or empty string>",
"trigger_words": ["<word1>", "<word2>"],
"short_description": "<1-2 sentence summary>",
"tags": ["<tag1>", "<tag2>"],
"recommended_width": 768,
"recommended_height": 1024,
"preview_url": "<image URL or empty string>",
"notes": "<plain-text usage summary or empty string>",
"usage_tips": "<JSON string like '{\"strength_min\":0.85,\"strength_max\":1.4}' or '{}'>",
"confidence": "<high|medium|low>"
}
```
Important:
- Only include the JSON object, no other text
- If a field cannot be determined, use an empty string or empty array
- Do not fabricate information not supported by the README
- Never use placeholder values like `"None"` or `"unknown"` for missing data — use empty string or empty array
File diff suppressed because it is too large Load Diff
+79 -2
View File
@@ -84,6 +84,7 @@ class Aria2Downloader:
self._transfers: Dict[str, Aria2Transfer] = {}
self._poll_interval = 0.5
self._state_store = Aria2TransferStateStore()
self._stderr_reader_task: Optional[asyncio.Task] = None
@property
def is_running(self) -> bool:
@@ -115,7 +116,7 @@ class Aria2Downloader:
try:
while True:
status = await self.get_status(download_id)
status = await self._get_status_with_retry(download_id)
if status is None:
return False, "aria2 download not found"
@@ -136,6 +137,35 @@ class Aria2Downloader:
finally:
self._transfers.pop(download_id, None)
async def _get_status_with_retry(
self, download_id: str, *, max_retries: int = 4, retry_delay: float = 3.0
) -> Optional[Dict[str, Any]]:
"""Call get_status with retry for transient RPC failures.
Only retries on :exc:`Aria2Error` (RPC-level failure). Returns
``None`` immediately when the download_id is not tracked (a missing
transfer is not a transient condition, so retrying is pointless).
A single failed RPC call should not immediately fail the download,
because aria2 may be temporarily busy (e.g. finalizing multiple
concurrent downloads) and a retry will often succeed.
"""
last_exc: Optional[Exception] = None
for attempt in range(max_retries):
try:
return await self.get_status(download_id)
except Aria2Error as exc:
last_exc = exc
if attempt < max_retries - 1:
logger.warning(
"aria2 get_status transient failure (attempt %d/%d) for %s: %s",
attempt + 1, max_retries, download_id, exc,
)
await asyncio.sleep(retry_delay)
raise Aria2Error(
f"Failed to query aria2 download status after {max_retries} attempts: {last_exc}"
) from last_exc
async def _schedule_download(
self,
url: str,
@@ -171,6 +201,13 @@ class Aria2Downloader:
"auto-file-renaming": "false",
"file-allocation": "none",
}
# Pass proxy to aria2 so the actual file transfer goes through the
# same proxy used by the aiohttp-based URL resolution step above.
downloader = await get_downloader()
if downloader.proxy_url:
options["all-proxy"] = downloader.proxy_url
if request_headers:
options["header"] = [
f"{key}: {value}" for key, value in request_headers.items()
@@ -312,6 +349,16 @@ class Aria2Downloader:
async def close(self) -> None:
"""Shut down the RPC process and session."""
# Cancel the background stderr reader first so it stops reading
# from the pipe before the subprocess is terminated.
if self._stderr_reader_task is not None:
self._stderr_reader_task.cancel()
try:
await asyncio.wait_for(self._stderr_reader_task, timeout=2.0)
except (asyncio.CancelledError, asyncio.TimeoutError):
pass
self._stderr_reader_task = None
if self._rpc_session is not None:
await self._rpc_session.close()
self._rpc_session = None
@@ -331,6 +378,23 @@ class Aria2Downloader:
process.kill()
await process.wait()
async def _drain_stderr(self) -> None:
"""Continuously drain aria2's stderr pipe so it never blocks.
When the 64 KB pipe buffer fills up, aria2's ``write()`` to stderr
blocks, which freezes the entire ``aria2c`` process including its
RPC handler. This background task reads lines from stderr as they
arrive and forwards them to Python's logger.
"""
try:
assert self._process is not None and self._process.stderr is not None
async for line in self._process.stderr:
text = line.decode("utf-8", errors="replace").rstrip()
if text:
logger.debug("aria2 stderr: %s", text)
except Exception:
pass
async def _dispatch_progress(self, callback, snapshot: DownloadProgress) -> None:
try:
result = callback(snapshot, snapshot)
@@ -465,6 +529,17 @@ class Aria2Downloader:
await self._wait_until_ready()
# Drain aria2's stderr in a background task so the pipe buffer
# never fills up. If the pipe blocks, aria2 itself freezes and
# cannot respond to RPC — this was the root cause of the
# "Failed to query aria2 download status" timeout bug.
# Must start AFTER _wait_until_ready to avoid a race where the
# drain task consumes aria2's early-exit error message before
# _wait_until_ready can read it.
self._stderr_reader_task = asyncio.create_task(
self._drain_stderr()
)
def _resolve_executable(self) -> str:
settings = get_settings_manager()
configured_path = (settings.get("aria2c_path") or "").strip()
@@ -584,7 +659,9 @@ class Aria2Downloader:
if self._rpc_session is None or self._rpc_session.closed:
async with self._rpc_session_lock:
if self._rpc_session is None or self._rpc_session.closed:
timeout = aiohttp.ClientTimeout(total=30)
timeout = aiohttp.ClientTimeout(
total=None, sock_connect=10, sock_read=60
)
self._rpc_session = aiohttp.ClientSession(timeout=timeout)
return self._rpc_session
+176 -9
View File
@@ -1,7 +1,7 @@
from abc import ABC, abstractmethod
import asyncio
import re
from typing import Any, Dict, List, Optional, Type, TYPE_CHECKING
from typing import Any, Dict, List, Optional, Type, Union, TYPE_CHECKING
import logging
import os
import time
@@ -104,6 +104,109 @@ class BaseModelService(ABC):
fetch_duration = time.perf_counter() - t0
initial_count = len(sorted_data)
# Optionally filter by civitai model ID (shows all local versions of a specific model)
civitai_model_id = kwargs.get("civitai_model_id")
if civitai_model_id is not None:
sorted_data = [
item for item in sorted_data
if self._extract_group_key(item) == civitai_model_id
]
# VLM mode: always sort by version ID descending (newest version first),
# regardless of the current sort_by preference.
# Fall back to modified timestamp for non-CivitAI sources.
sorted_data.sort(
key=lambda x: self._extract_version_id(x)
or x.get("modified", 0)
or 0,
reverse=True,
)
# Optionally group by civitai modelId, showing only the latest version per model
dedup_lost = 0
if kwargs.get("group_by_model") and civitai_model_id is None:
# Determine whether to further sub-group by base model
# When version_grouping is "same_base", versions with different
# base models are effectively different groups — the dedup key
# needs to include base_model so the version count and VLM flow
# stay consistent (card shows correct count for its base model).
ufs = self.settings.get("version_grouping", "same_base")
group_by_base = ufs == "same_base"
dedup_map = {} # (modelId [,base_model]) -> (item, version_or_modified)
version_counter = {} # same-key -> count
standalone = []
for item in sorted_data:
mid = self._extract_group_key(item)
if mid is None:
standalone.append(item)
continue
key = (mid, item.get("base_model") or "") if group_by_base else mid
# Count all versions per key
version_counter[key] = version_counter.get(key, 0) + 1
# Prefer CivitAI version_id; fall back to modified timestamp
vid = self._extract_version_id(item)
if vid is None:
vid = item.get("modified", 0) or 0
if key not in dedup_map or vid > dedup_map[key][1]:
dedup_map[key] = (item, vid)
# Attach version_count to each surviving grouped item (shallow copy
# to avoid mutating cached dicts — the cache is shared across requests)
for key, (item, vid) in dedup_map.items():
item = dict(item)
item["version_count"] = version_counter[key]
dedup_map[key] = (item, vid)
dedup_lost = len(sorted_data) - (len(dedup_map) + len(standalone))
sorted_data = [entry[0] for entry in dedup_map.values()] + standalone
# Re-sort by version_count (grouped: after dedup; non-grouped: group internally, sort, expand)
if sort_params.key == "versions_count" and civitai_model_id is None:
reverse = sort_params.order == "desc"
if kwargs.get("group_by_model"):
# Grouped mode: items are already dedup'd with version_count attached
sorted_data.sort(
key=lambda x: (
x.get("version_count", 0),
(x.get("model_name") or x.get("file_name") or "").lower(),
x.get("file_path", "").lower(),
),
reverse=reverse,
)
else:
# Non-grouped mode: group internally, sort groups by count, expand
# Respect the version_grouping setting (same logic as grouped dedup)
ufs = self.settings.get("version_grouping", "same_base")
group_by_base = ufs == "same_base"
model_groups: Dict[Any, List[Dict]] = {}
ungrouped_standalone: List[Dict] = []
for item in sorted_data:
mid = self._extract_group_key(item)
if mid is None:
ungrouped_standalone.append(item)
continue
key = (mid, item.get("base_model") or "") if group_by_base else mid
model_groups.setdefault(key, []).append(item)
# Sort versions within each group by version id (descending);
# fall back to modified timestamp for non-CivitAI sources.
for items in model_groups.values():
items.sort(
key=lambda x: self._extract_version_id(x)
or x.get("modified", 0)
or 0,
reverse=True,
)
# Sort groups by version count
sorted_groups = sorted(
model_groups.values(),
key=lambda items: len(items),
reverse=reverse,
)
# Flatten: grouped items first, standalone items last
sorted_data = []
for items in sorted_groups:
sorted_data.extend(items)
sorted_data.extend(ungrouped_standalone)
t1 = time.perf_counter()
if hash_filters:
filtered_data = await self._apply_hash_filters(sorted_data, hash_filters)
@@ -172,7 +275,7 @@ class BaseModelService(ABC):
overall_duration = time.perf_counter() - overall_start
logger.debug(
"%s.get_paginated_data took %.3fs (fetch: %.3fs, filter: %.3fs, update_filter: %.3fs, pagination: %.3fs, annotate: %.3fs). "
"Counts: initial=%d, post_filter=%d, final=%d",
"Counts: initial=%d, dedup=%d, post_filter=%d, final=%d",
self.__class__.__name__,
overall_duration,
fetch_duration,
@@ -181,6 +284,7 @@ class BaseModelService(ABC):
pagination_duration,
annotate_duration,
initial_count,
dedup_lost,
post_filter_count,
final_count,
)
@@ -495,7 +599,7 @@ class BaseModelService(ABC):
if not ordered_ids:
return annotated
strategy_value = self.settings.get("update_flag_strategy")
strategy_value = self.settings.get("version_grouping")
if isinstance(strategy_value, str) and strategy_value.strip():
strategy = strategy_value.strip().lower()
else:
@@ -602,6 +706,33 @@ class BaseModelService(ABC):
return annotated
@staticmethod
def _extract_hf_group_key(item: Dict) -> Optional[str]:
"""Extract `hf:{owner}/{repo}` from item's ``hf_url``, or None."""
hf_url = item.get("hf_url") if isinstance(item, dict) else None
if not hf_url or not isinstance(hf_url, str):
return None
m = re.match(
r"https?://huggingface\.co/([^/]+/[^/]+)", hf_url.strip()
)
if not m:
return None
return f"hf:{m.group(1)}"
@staticmethod
def _extract_group_key(item: Dict) -> Union[int, str, None]:
"""Return the group identity key: CivitAI modelId (int) or HF repo (str).
Preference order:
1. CivitAI ``modelId`` (int)
2. HF repo identity ``hf:{owner}/{repo}`` (str)
3. ``None`` (no known grouping source)
"""
mid = BaseModelService._extract_model_id(item)
if mid is not None:
return mid
return BaseModelService._extract_hf_group_key(item)
@staticmethod
def _extract_model_id(item: Dict) -> Optional[int]:
civitai = item.get("civitai") if isinstance(item, dict) else None
@@ -696,8 +827,12 @@ class BaseModelService(ABC):
}
@abstractmethod
async def format_response(self, model_data: Dict) -> Dict:
"""Format model data for API response - must be implemented by subclasses"""
async def format_response(self, model_data: Dict) -> Optional[Dict]:
"""Format model data for API response - must be implemented by subclasses.
Subclasses should return None for corrupted entries so the handler
layer can filter them out. See issue #730.
"""
pass
# Common service methods that delegate to scanner
@@ -705,6 +840,12 @@ class BaseModelService(ABC):
"""Get top tags sorted by frequency"""
return await self.scanner.get_top_tags(limit)
async def search_tags(
self, query: str, limit: int = 50
) -> List[Dict]:
"""Search tags by substring, sorted by frequency"""
return await self.scanner.search_tags(query, limit)
async def get_base_models(self, limit: int = 20) -> List[Dict]:
"""Get base models sorted by frequency"""
return await self.scanner.get_base_models(limit)
@@ -856,13 +997,21 @@ class BaseModelService(ABC):
return unified_tree
async def get_model_notes(self, model_name: str) -> Optional[str]:
"""Get notes for a specific model file"""
async def get_model_notes(self, model_name: str) -> Optional[dict]:
"""Get notes and file_path for a specific model file.
Supports both simple names (``OWSMianne_ANIMA_V1``) and full-path
syntax (``Anima/character/OWSMianne_ANIMA_V1``).
"""
cache = await self.scanner.get_cached_data()
for model in cache.raw_data:
if model["file_name"] == model_name:
return model.get("notes", "")
file_name = model.get("file_name", "")
if file_name == model_name or model_name.endswith("/" + file_name) or model_name.endswith("\\" + file_name):
return {
"notes": model.get("notes", ""),
"file_path": model.get("file_path", ""),
}
return None
@@ -985,6 +1134,11 @@ class BaseModelService(ABC):
Listing/search endpoints return lightweight cache entries; this method performs
a lazy read of the on-disk metadata snapshot when callers need full detail.
As a beneficial side effect, the in-memory and persistent caches are
opportunistically synchronised with the on-disk metadata this keeps the
caches fresh even when a ``.metadata.json`` file was edited outside of the
normal save path (e.g. manually or by an external script).
"""
metadata, should_skip = await MetadataManager.load_metadata(
file_path, self.metadata_class
@@ -1002,6 +1156,19 @@ class BaseModelService(ABC):
MetadataManager.save_metadata(file_path, metadata)
)
# Opportunistically sync the in-memory + persistent caches.
# The .metadata.json disk read is already paid for; the sync only
# performs work when the cache is actually stale, and uses targeted,
# in-place operations to minimise overhead even with large model sets.
#
# Fire-and-forget by design: the task is intentionally untracked.
# sync_cache_from_metadata handles its own errors internally.
asyncio.create_task(
self.scanner.sync_cache_from_metadata(
file_path, metadata.to_dict()
)
)
return self.filter_civitai_data(metadata.to_dict().get("civitai", {}))
async def get_model_description(self, file_path: str) -> Optional[str]:
+4
View File
@@ -523,6 +523,10 @@ class BatchImportService:
if payload.get("checkpoint"):
metadata["checkpoint"] = payload["checkpoint"]
nsfw = payload.get("preview_nsfw_level")
if isinstance(nsfw, int) and nsfw > 0:
metadata["preview_nsfw_level"] = nsfw
image_bytes = None
image_base64 = payload.get("image_base64")
+25
View File
@@ -114,6 +114,13 @@ class CheckpointScanner(ModelScanner):
and metadata.hash_status == "completed"
and metadata.sha256
):
# Ensure the in-memory hash index is populated even when
# the hash was already computed and persisted to the metadata
# file. Without this, usage tracking (and any other caller
# that queries get_hash_by_filename first) will miss on every
# lookup and keep calling back into this method, creating a
# tight loop that never populates the index.
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
return metadata.sha256
async with self._hash_calculation_lock:
@@ -125,6 +132,7 @@ class CheckpointScanner(ModelScanner):
and metadata.hash_status == "completed"
and metadata.sha256
):
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
return metadata.sha256
task = self._hash_calculation_tasks.get(real_path)
@@ -175,6 +183,9 @@ class CheckpointScanner(ModelScanner):
# Check if hash is already calculated
if metadata.hash_status == "completed" and metadata.sha256:
# Populate the in-memory hash index even for pre-computed
# hashes, mirroring the fix in calculate_hash_for_model.
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
return metadata.sha256
# Update status to calculating
@@ -193,6 +204,20 @@ class CheckpointScanner(ModelScanner):
# Update hash index
self._hash_index.add_entry(sha256.lower(), file_path)
# Update the in-memory cache entry so that subsequent
# _persist_current_cache / _save_persistent_cache calls
# write the hash back to the SQLite models table. Without
# this the hash only lives in the metadata file and the
# in-memory hash index, both of which are lost across
# restarts, causing the same re-computation loop on the
# next session.
if self._cache is not None and self._cache.raw_data:
for entry in self._cache.raw_data:
if entry.get("file_path") == file_path:
entry["sha256"] = sha256.lower()
entry["hash_status"] = "completed"
break
logger.info(f"Hash calculated for checkpoint: {file_path}")
return sha256
+27 -8
View File
@@ -1,6 +1,6 @@
import os
import logging
from typing import Dict
from typing import Dict, Optional
from .base_model_service import BaseModelService
from .auto_tag_service import extract_auto_tags
@@ -21,20 +21,37 @@ class CheckpointService(BaseModelService):
"""
super().__init__("checkpoint", scanner, CheckpointMetadata, update_service=update_service)
async def format_response(self, checkpoint_data: Dict) -> Dict:
"""Format Checkpoint data for API response"""
async def format_response(self, checkpoint_data: Dict) -> Optional[Dict]:
"""Format Checkpoint data for API response.
Returns None when the entry is missing critical fields (corrupted cache
row), so the handler layer can filter it out. See issue #730.
"""
# Guard against corrupted cache entries missing critical fields
file_path = checkpoint_data.get("file_path")
if not file_path or not isinstance(file_path, str):
logger.warning(
"Skipping corrupted checkpoint entry (missing file_path): %s",
checkpoint_data.get("file_name", "<unknown>"),
)
return None
# Get sub_type from cache entry (new canonical field)
sub_type = checkpoint_data.get("sub_type", "checkpoint")
file_name = checkpoint_data.get("file_name") or ""
model_name = checkpoint_data.get("model_name") or file_name
folder = checkpoint_data.get("folder") or ""
return {
"model_name": checkpoint_data["model_name"],
"file_name": checkpoint_data["file_name"],
"model_name": model_name,
"file_name": file_name,
"preview_url": config.get_preview_static_url(checkpoint_data.get("preview_url", "")),
"preview_nsfw_level": checkpoint_data.get("preview_nsfw_level", 0),
"base_model": checkpoint_data.get("base_model", ""),
"folder": checkpoint_data["folder"],
"folder": folder,
"sha256": checkpoint_data.get("sha256", ""),
"file_path": checkpoint_data["file_path"].replace(os.sep, "/"),
"file_path": file_path.replace(os.sep, "/"),
"file_size": checkpoint_data.get("size", 0),
"modified": checkpoint_data.get("modified", ""),
"tags": checkpoint_data.get("tags", []),
@@ -48,6 +65,8 @@ class CheckpointService(BaseModelService):
"skip_metadata_refresh": bool(checkpoint_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(checkpoint_data.get("civitai", {}), minimal=True),
"auto_tags": checkpoint_data.get("auto_tags") or extract_auto_tags(checkpoint_data),
"version_count": checkpoint_data.get("version_count"),
"hf_url": checkpoint_data.get("hf_url", ""),
}
def find_duplicate_hashes(self) -> Dict:
+16 -2
View File
@@ -304,6 +304,20 @@ class CivArchiveClient:
version_id = file_data.get("model_version_id") or file_data.get("modelVersionId")
if model_id is None or version_id is None:
continue
# CivitAI / CivArchive model IDs are small integers (typically ≤ 7
# digits). Reject suspiciously large values that indicate the API
# returned a malformed payload (e.g. a hash reinterpreted as an ID)
# to avoid pointless HTTP 500 errors from CivArchive.
_MAX_VALID_CIVITAI_ID = 100_000_000
try:
if int(model_id) >= _MAX_VALID_CIVITAI_ID or int(version_id) >= _MAX_VALID_CIVITAI_ID:
logger.debug(
"Skipping implausible CivArchive model_id=%s / version_id=%s",
model_id, version_id,
)
continue
except (TypeError, ValueError):
continue
resolved = await self.get_model_version(model_id, version_id)
if resolved:
return resolved
@@ -327,7 +341,7 @@ class CivArchiveClient:
if resolved:
return resolved, None
logger.error("Error fetching version of CivArchive model by hash %s", model_hash[:10])
logger.debug("Error fetching version of CivArchive model by hash %s", model_hash[:10])
return None, "No version data found"
except RateLimitError:
@@ -417,7 +431,7 @@ class CivArchiveClient:
if version_id is not None:
raw_id = version_data.get("id")
if raw_id != version_id:
if raw_id is not None and str(raw_id) != str(version_id):
logger.warning(
"Requested version %s doesn't match default version %s for model %s",
version_id,
+26
View File
@@ -196,6 +196,7 @@ class CivitaiBaseModelService:
"ernie": "ERNI",
"ernie turbo": "ETRB",
"nucleus": "NUCL",
"krea 2": "KR2",
"svd": "SVD",
"ltxv": "LTXV",
"ltxv2": "LTV2",
@@ -212,6 +213,18 @@ class CivitaiBaseModelService:
"wan video 2.2 i2v-a14b": "WAN",
"wan video 2.5 t2v": "WAN",
"wan video 2.5 i2v": "WAN",
"wan video 2.7": "WAN",
"wan image 2.7": "WI27",
"ace audio": "ACE",
"boogu": "BOOG",
"grok": "GROK",
"happyhorse": "HAPP",
"hidream-o1": "HIO1",
"lens": "LENS",
"mai": "MAI",
"upscaler": "UPSC",
"ideogram 4.0": "ID40",
"qwen 2": "QWN2",
}
if lower_name in special_cases:
@@ -391,6 +404,7 @@ class CivitaiBaseModelService:
"LTXV2",
"LTXV 2.3",
"CogVideoX",
"HappyHorse",
"Mochi",
"Hunyuan Video",
"Wan Video",
@@ -403,15 +417,25 @@ class CivitaiBaseModelService:
"Wan Video 2.2 I2V-A14B",
"Wan Video 2.5 T2V",
"Wan Video 2.5 I2V",
"Wan Image 2.7",
"Wan Video 2.7",
],
"Other Models": [
"ACE Audio",
"Illustrious",
"Pony",
"Pony V7",
"Boogu",
"HiDream",
"HiDream-O1",
"Ideogram 4.0",
"Qwen",
"Qwen 2",
"AuraFlow",
"Chroma",
"Grok",
"Lens",
"MAI",
"ZImageTurbo",
"ZImageBase",
"PixArt a",
@@ -424,6 +448,8 @@ class CivitaiBaseModelService:
"Ernie",
"Ernie Turbo",
"Nucleus",
"Krea 2",
"Upscaler",
],
}
+1 -1
View File
@@ -56,7 +56,7 @@ class CivitaiClient:
self._MAX_CACHE_ENTRIES = 500
def _build_image_info_url(self, image_id: str) -> str:
return f"{self.base_url}/images?imageId={image_id}&nsfw=X"
return f"{self.base_url}/images?imageId={image_id}&nsfw=X&withMeta=true"
async def _make_request(
self,
+177 -48
View File
@@ -29,6 +29,7 @@ from .metadata_service import get_default_metadata_provider, get_metadata_provid
from .downloader import get_downloader, DownloadProgress, DownloadStreamControl
from .aria2_downloader import Aria2Error, get_aria2_downloader
from .aria2_transfer_state import Aria2TransferStateStore
from .download_queue_service import DownloadQueueService
# Download to temporary file first
import tempfile
@@ -229,6 +230,12 @@ class DownloadManager:
Returns:
Dict with download result
"""
logger.debug(
"[download] download_from_civitai called: model_id=%s, model_version_id=%s, "
"source=%s, file_params=%s",
model_id, model_version_id, source, file_params,
)
# Validate that at least one identifier is provided
if not model_id and not model_version_id:
return {
@@ -249,6 +256,7 @@ class DownloadManager:
"source": source,
"file_params": copy.deepcopy(file_params) if file_params is not None else None,
"progress": 0,
"status": "queued",
"transfer_backend": self._get_model_download_backend(),
"bytes_downloaded": 0,
@@ -288,8 +296,8 @@ class DownloadManager:
return result
except asyncio.CancelledError:
return {
"success": False,
"error": "Download was cancelled",
"success": True,
"cancelled": True,
"download_id": task_id,
}
finally:
@@ -360,6 +368,15 @@ class DownloadManager:
if self._active_downloads[task_id].get("transfer_backend") == "aria2":
await self._persist_aria2_state(task_id)
# Update SQLite queue status to 'downloading'
try:
queue_service = await DownloadQueueService.get_instance()
await queue_service.update_status(task_id, "downloading")
except Exception:
logger.warning(
"Failed to update queue status for %s", task_id, exc_info=True
)
# Use original download implementation
try:
# Check for cancellation before starting
@@ -396,6 +413,22 @@ class DownloadManager:
if self._active_downloads[task_id].get("transfer_backend") == "aria2":
await self._persist_aria2_state(task_id)
# Move queue item to history on completion
try:
queue_service = await DownloadQueueService.get_instance()
await queue_service.complete_download(
download_id=task_id,
status=result.get("status", "completed") if result.get("success") else "failed",
error=result.get("error") if not result.get("success") else None,
file_path=result.get("file_path"),
bytes_downloaded=self._active_downloads.get(task_id, {}).get("bytes_downloaded", 0),
total_bytes=self._active_downloads.get(task_id, {}).get("total_bytes"),
)
except Exception:
logger.warning(
"Failed to complete queue item for %s", task_id, exc_info=True
)
return result
except asyncio.CancelledError:
# Handle cancellation
@@ -404,6 +437,19 @@ class DownloadManager:
self._active_downloads[task_id]["bytes_per_second"] = 0.0
if self._active_downloads[task_id].get("transfer_backend") == "aria2":
await self._persist_aria2_state(task_id)
# Move queue item to history as canceled
try:
queue_service = await DownloadQueueService.get_instance()
await queue_service.complete_download(
download_id=task_id,
status="canceled",
)
except Exception:
logger.warning(
"Failed to cancel queue item for %s", task_id, exc_info=True
)
logger.info(f"Download cancelled for task {task_id}")
raise
except Exception as e:
@@ -417,6 +463,22 @@ class DownloadManager:
self._active_downloads[task_id]["bytes_per_second"] = 0.0
if self._active_downloads[task_id].get("transfer_backend") == "aria2":
await self._persist_aria2_state(task_id)
# Move queue item to history as failed
try:
queue_service = await DownloadQueueService.get_instance()
await queue_service.complete_download(
download_id=task_id,
status="failed",
error=str(e),
bytes_downloaded=self._active_downloads.get(task_id, {}).get("bytes_downloaded", 0),
total_bytes=self._active_downloads.get(task_id, {}).get("total_bytes"),
)
except Exception:
logger.warning(
"Failed to complete queue item for %s", task_id, exc_info=True
)
return {"success": False, "error": str(e)}
finally:
# Schedule cleanup of download record after delay
@@ -620,7 +682,10 @@ class DownloadManager:
u for u in download_urls if not u.startswith(CIVITAI_DOWNLOAD_URL_PREFIXES)
]
download_urls = non_civitai_urls + civitai_urls
else:
# Fallback: when mirrors is empty or all mirrors have been deleted,
# use the file's downloadUrl directly (e.g. CivitAI download endpoint).
if not download_urls:
download_url = file_info.get("downloadUrl")
if download_url:
download_urls.append(normalize_civitai_download_url(download_url))
@@ -1233,10 +1298,24 @@ class DownloadManager:
"download_id": download_id,
}
# Check if this checkpoint should be treated as a diffusion model based on baseModel
# Check if this checkpoint should be treated as a diffusion model
# Priority: (1) any file has type "UNet" or "Diffusion Model",
# (2) baseModel is in DIFFUSION_MODEL_BASE_MODELS
is_diffusion_model = False
if model_type == "checkpoint":
if base_model_value in DIFFUSION_MODEL_BASE_MODELS:
# Check file types first (more direct signal from CivitAI)
version_files = version_info.get("files", [])
for f in version_files:
f_type = f.get("type", "")
if f_type in ("UNet", "Diffusion Model"):
is_diffusion_model = True
logger.info(
f"File type '{f_type}' detected, routing checkpoint to unet folder"
)
break
# Fallback to baseModel name check
if not is_diffusion_model and base_model_value in DIFFUSION_MODEL_BASE_MODELS:
is_diffusion_model = True
logger.info(
f"baseModel '{base_model_value}' is a known diffusion model, routing to unet folder"
@@ -1310,7 +1389,17 @@ class DownloadManager:
# Update save directory with relative path if provided
if relative_path:
base_save_dir = save_dir
save_dir = os.path.join(save_dir, relative_path)
# Security: validate path containment after joining
resolved_dir = os.path.abspath(os.path.normpath(save_dir))
base_dir = os.path.abspath(os.path.normpath(base_save_dir))
if not resolved_dir.startswith(base_dir + os.sep) and resolved_dir != base_dir:
logger.warning(
"Path traversal detected: %s escapes %s",
resolved_dir, base_dir,
)
return {"success": False, "error": "Download path is outside allowed directory"}
# Create directory if it doesn't exist
os.makedirs(save_dir, exist_ok=True)
@@ -1352,86 +1441,100 @@ class DownloadManager:
# If file_params is provided, try to find matching file
if file_params and model_version_id:
target_file_id = file_params.get("id")
target_type = file_params.get("type", "Model")
target_format = file_params.get("format", "SafeTensor")
target_size = file_params.get("size", "full")
target_format = file_params.get("format")
target_size = file_params.get("size")
target_fp = file_params.get("fp")
is_primary = file_params.get("isPrimary", False)
if is_primary:
# Find primary file
logger.debug(
"[download] file_params received: id=%s, type=%s, format=%s, size=%s, fp=%s, isPrimary=%s, "
"model_version_id=%s, total_files=%d",
target_file_id, target_type, target_format, target_size, target_fp, is_primary,
model_version_id, len(files),
)
if target_file_id:
target_id_str = str(target_file_id)
for f in files:
f_id = f.get("id")
if str(f_id) == target_id_str:
file_info = f
logger.debug(
"[download] MATCH by ID: id=%s name='%s'",
f_id, f.get("name"),
)
break
if not file_info:
logger.debug("[download] No file found with id=%s", target_file_id)
elif is_primary:
file_info = next(
(
f
for f in files
if f.get("primary")
and f.get("type") in ("Model", "Negative", "Diffusion Model")
and f.get("type") in ("Model", "Negative", "Diffusion Model", "UNet")
),
None,
)
else:
# Match by metadata
# Lenient metadata match: only compare fields present on both sides
for f in files:
f_type = f.get("type", "")
f_meta = f.get("metadata", {})
# Check type match
if f_type != target_type:
continue
# Check metadata match
if f_meta.get("format") != target_format:
f_meta = f.get("metadata", {})
f_format = f_meta.get("format") or f.get("format")
f_size = f_meta.get("size") or f.get("size")
f_fp = f_meta.get("fp") or f.get("fp")
if target_format and f_format != target_format:
continue
if f_meta.get("size") != target_size:
if target_size and f_size and f_size != target_size:
continue
if target_fp and f_meta.get("fp") != target_fp:
if target_fp and f_fp and f_fp != target_fp:
continue
file_info = f
break
if not file_info:
logger.debug(
"[download] No match found via file_params — falling back to primary file lookup",
)
elif not file_params:
logger.debug(
"[download] No file_params provided (null/None) — will use primary file lookup. "
"model_version_id=%s, total_files=%d",
model_version_id, len(files),
)
# Fallback to primary file if no match found
if not file_info:
logger.debug("[download] Looking for primary file as fallback")
file_info = next(
(
f
for f in files
if f.get("primary") and f.get("type") in ("Model", "Negative", "Diffusion Model")
if f.get("primary") and f.get("type") in ("Model", "Negative", "Diffusion Model", "UNet")
),
None,
)
if file_info:
logger.debug(
"[download] Fallback primary file selected: id=%s, name=%s",
file_info.get("id"), file_info.get("name"),
)
else:
logger.debug("[download] No primary file found in fallback lookup")
if not file_info:
return {"success": False, "error": "No suitable file found in metadata"}
mirrors = file_info.get("mirrors") or []
download_urls = []
if mirrors:
for mirror in mirrors:
if mirror.get("deletedAt") is None and mirror.get("url"):
download_urls.append(
normalize_civitai_download_url(mirror["url"])
)
# When source is 'civarchive', prioritize non-Civitai URLs
# This avoids failed downloads from deleted Civitai models
if source == "civarchive" and len(download_urls) > 1:
civitai_urls = [
u
for u in download_urls
if u.startswith(CIVITAI_DOWNLOAD_URL_PREFIXES)
]
non_civitai_urls = [
u
for u in download_urls
if not u.startswith(CIVITAI_DOWNLOAD_URL_PREFIXES)
]
download_urls = non_civitai_urls + civitai_urls
else:
download_url = file_info.get("downloadUrl")
if download_url:
download_urls.append(
normalize_civitai_download_url(download_url)
)
download_urls = self._build_download_urls_from_file_info(file_info, source=source)
if not download_urls:
return {"success": False, "error": "No mirror URL found"}
@@ -1734,6 +1837,9 @@ class DownloadManager:
model_tags, model_type
)
if not first_tag:
first_tag = "no tags" # Default if no tags available
# Format the template with available data
formatted_path = path_template
formatted_path = formatted_path.replace("{base_model}", mapped_base_model)
@@ -1749,6 +1855,15 @@ class DownloadManager:
if model_type == "embedding":
formatted_path = formatted_path.replace(" ", "_")
# Sanitize the resolved path to prevent path traversal:
# - Strip leading slashes (prevents os.path.join from treating path as absolute)
# - Collapse double slashes from empty placeholder substitutions
# - Strip trailing slashes for cleanliness
formatted_path = formatted_path.lstrip("/")
while "//" in formatted_path:
formatted_path = formatted_path.replace("//", "/")
formatted_path = formatted_path.rstrip("/")
return formatted_path
async def _execute_download(
@@ -1974,7 +2089,21 @@ class DownloadManager:
break
last_error = result
if os.path.exists(save_path):
# For aria2: if the .aria2 control file is missing, aria2 considers
# the download complete. A transient RPC failure may have made us
# think the download failed even though the file is fully on disk.
# Keep the file so a retry can find it already complete.
if (
transfer_backend == "aria2"
and os.path.exists(save_path)
and not os.path.exists(f"{save_path}.aria2")
):
logger.warning(
"aria2 download reported failure but .aria2 file is absent "
"for %s — the file is likely complete. Preserving it for retry.",
save_path,
)
elif os.path.exists(save_path):
try:
os.remove(save_path)
except Exception as e:
+88 -21
View File
@@ -31,7 +31,7 @@ class DownloadQueueService:
_instance: Optional[DownloadQueueService] = None
_class_lock: asyncio.Lock = asyncio.Lock()
_SCHEMA = """
_SCHEMA_TABLES = """
CREATE TABLE IF NOT EXISTS download_queue (
download_id TEXT PRIMARY KEY,
model_id INTEGER,
@@ -76,6 +76,11 @@ class DownloadQueueService:
CREATE INDEX IF NOT EXISTS idx_dh_status ON download_history(status);
"""
_CREATE_UNIQUE_INDEX = """
CREATE UNIQUE INDEX IF NOT EXISTS idx_dh_download_id
ON download_history(download_id) WHERE download_id IS NOT NULL;
"""
@classmethod
async def get_instance(cls) -> DownloadQueueService:
"""Return the singleton instance, creating it if necessary."""
@@ -113,10 +118,39 @@ class DownloadQueueService:
if self._schema_initialized:
return
with self._connect() as conn:
conn.executescript(self._SCHEMA)
conn.executescript(self._SCHEMA_TABLES)
# Creating the unique index on download_history.download_id can
# fail if pre-existing rows have duplicate values (e.g. from a
# previous version that lacked the index). Deduplicate first so
# that the migration does not crash on startup.
if not self._index_exists(conn, "idx_dh_download_id"):
self._remove_duplicate_download_ids(conn)
conn.executescript(self._CREATE_UNIQUE_INDEX)
conn.commit()
self._schema_initialized = True
@staticmethod
def _index_exists(conn: sqlite3.Connection, name: str) -> bool:
return conn.execute(
"SELECT 1 FROM sqlite_master WHERE type='index' AND name=?",
(name,),
).fetchone() is not None
@staticmethod
def _remove_duplicate_download_ids(conn: sqlite3.Connection) -> None:
conn.execute("""
DELETE FROM download_history
WHERE id NOT IN (
SELECT MIN(id)
FROM download_history
WHERE download_id IS NOT NULL
GROUP BY download_id
)
AND download_id IS NOT NULL
""")
def get_database_path(self) -> str:
"""Return the resolved database file path."""
return self._db_path
@@ -154,13 +188,23 @@ class DownloadQueueService:
"""Insert a new download into the queue.
Returns the inserted row as a dict (or an empty dict if the
download_id already exists).
download_id already exists in the queue or has a terminal
record in history).
"""
now = time.time()
file_params_json = json.dumps(file_params) if file_params is not None else None
async with self._lock:
conn = self._get_conn()
# Reject download_ids that already have a terminal record in history.
history_row = conn.execute(
"SELECT 1 FROM download_history WHERE download_id = ? LIMIT 1",
(download_id,),
).fetchone()
if history_row is not None:
return {}
conn.execute(
"""
INSERT OR IGNORE INTO download_queue (
@@ -380,7 +424,7 @@ class DownloadQueueService:
)
conn.execute(
"""
INSERT INTO download_history (
INSERT OR IGNORE INTO download_history (
download_id, model_id, model_version_id, model_name,
version_name, thumbnail_url, status, error, file_path,
bytes_downloaded, total_bytes, completed_at
@@ -537,17 +581,27 @@ class DownloadQueueService:
"offset": offset,
}
async def delete_history_item(self, id: int) -> bool:
"""Delete a single history entry by its *id*.
async def delete_history_item(
self, id: Optional[int] = None, download_id: Optional[str] = None
) -> bool:
"""Delete a single history entry by *download_id* (preferred) or *id*.
Returns ``True`` if a row was deleted.
"""
async with self._lock:
conn = self._get_conn()
cursor = conn.execute(
"DELETE FROM download_history WHERE id = ?",
(id,),
)
if download_id:
cursor = conn.execute(
"DELETE FROM download_history WHERE download_id = ?",
(download_id,),
)
elif id is not None:
cursor = conn.execute(
"DELETE FROM download_history WHERE id = ?",
(id,),
)
else:
return False
conn.commit()
return cursor.rowcount > 0
@@ -604,21 +658,34 @@ class DownloadQueueService:
# Retry
# ------------------------------------------------------------------
async def retry_from_history(self, item_id: int) -> Optional[dict[str, Any]]:
async def retry_from_history(
self,
item_id: Optional[int] = None,
download_id: Optional[str] = None,
) -> Optional[dict[str, Any]]:
"""Re-queue a failed or canceled download from history.
Looks up the history record by its primary key. If the status is
``failed`` or ``canceled`` a new queue entry is created with the
same model metadata and a fresh download id, and the original
history entry is **deleted** to prevent exponential growth when
the retried item is later canceled or fails again and re-retried.
Looks up the history record by *download_id* (preferred) or
*item_id*. If the status is ``failed`` or ``canceled`` a new
queue entry is created with the same model metadata and a fresh
download id, and the original history entry is **deleted** to
prevent exponential growth when the retried item is later
canceled or fails again and re-retried.
"""
async with self._lock:
conn = self._get_conn()
row = conn.execute(
"SELECT * FROM download_history WHERE id = ?",
(item_id,),
).fetchone()
if download_id:
row = conn.execute(
"SELECT * FROM download_history WHERE download_id = ?",
(download_id,),
).fetchone()
elif item_id is not None:
row = conn.execute(
"SELECT * FROM download_history WHERE id = ?",
(item_id,),
).fetchone()
else:
return None
if row is None:
return None
status = str(row["status"])
@@ -650,7 +717,7 @@ class DownloadQueueService:
)
conn.execute(
"DELETE FROM download_history WHERE id = ?",
(item_id,),
(row["id"],),
)
conn.commit()
queued = conn.execute(
+56 -10
View File
@@ -46,6 +46,30 @@ def is_ssl_cert_verify_error(exc: BaseException) -> bool:
return "CERTIFICATE_VERIFY_FAILED" in str(exc)
def _parse_retry_after(value: str) -> int:
"""Parse a Retry-After header value into seconds.
Supports both integer seconds and HTTP-date formats.
Returns a default of 60 seconds on invalid/missing input.
"""
if not value or not value.strip():
return 60
value = value.strip()
try:
return max(1, int(value))
except ValueError:
pass
try:
parsed = parsedate_to_datetime(value)
now = datetime.now().astimezone()
delta = (parsed - now).total_seconds()
return max(1, int(delta))
except (ValueError, OverflowError, OSError):
return 60
@dataclass(frozen=True)
class DownloadProgress:
"""Snapshot of a download transfer at a moment in time."""
@@ -246,14 +270,14 @@ class Downloader:
Note: This is private and caller MUST hold self._session_lock.
"""
# Close existing session if any
if self._session is not None:
try:
await self._session.close()
except Exception as e: # pragma: no cover
logger.warning(f"Error closing previous session: {e}")
finally:
self._session = None
# Snapshot and clear old session reference before creating the new
# one. This ensures self._session is always valid (or None, which
# triggers a fresh creation) and avoids a race where concurrent
# requests hold a reference to a session whose connector has been
# torn down by a premature close() call — the root cause of the
# intermittent "NoneType has no attribute connect" crash.
old_session = self._session
self._session = None
# Check for app-level proxy settings
proxy_url = None # http(s) proxy, passed via the per-request `proxy=` kwarg
@@ -348,6 +372,13 @@ class Downloader:
self._proxy_url = proxy_url
self._session_created_at = datetime.now()
# Close the previous session now that the replacement is live.
if old_session is not None:
try:
await old_session.close()
except Exception as e: # pragma: no cover
logger.warning(f"Error closing previous session: {e}")
logger.debug(
"Created new HTTP session with proxy settings. App-level proxy: %s, System-level proxy (trust_env): %s",
bool(proxy_url),
@@ -729,7 +760,8 @@ class Downloader:
else:
resume_offset = 0
total_size = 0
await self._create_session()
async with self._session_lock:
await self._create_session()
continue
return False, integrity_error
@@ -819,7 +851,8 @@ class Downloader:
logger.info(f"Will resume from byte {resume_offset}")
# Refresh session to get new connection
await self._create_session()
async with self._session_lock:
await self._create_session()
continue
else:
logger.error(f"Max retries exceeded for download: {e}")
@@ -911,6 +944,19 @@ class Downloader:
elif response.status == 404:
error_msg = "File not found"
return False, error_msg, None
elif response.status == 429:
raw_retry_after = response.headers.get("Retry-After")
retry_after = _parse_retry_after(raw_retry_after or "")
if raw_retry_after:
logger.warning(
"Rate limited (429) for %s, Retry-After: %ss", url, retry_after
)
else:
logger.warning(
"Rate limited (429) for %s, no Retry-After header; defaulting to %ss",
url, retry_after,
)
return False, f"Rate limited (429), retry after {retry_after}s", None
else:
error_msg = f"Download failed with status {response.status}"
return False, error_msg, None
+27 -8
View File
@@ -1,6 +1,6 @@
import os
import logging
from typing import Dict
from typing import Dict, Optional
from .base_model_service import BaseModelService
from .auto_tag_service import extract_auto_tags
@@ -21,20 +21,37 @@ class EmbeddingService(BaseModelService):
"""
super().__init__("embedding", scanner, EmbeddingMetadata, update_service=update_service)
async def format_response(self, embedding_data: Dict) -> Dict:
"""Format Embedding data for API response"""
async def format_response(self, embedding_data: Dict) -> Optional[Dict]:
"""Format Embedding data for API response.
Returns None when the entry is missing critical fields (corrupted cache
row), so the handler layer can filter it out. See issue #730.
"""
# Guard against corrupted cache entries missing critical fields
file_path = embedding_data.get("file_path")
if not file_path or not isinstance(file_path, str):
logger.warning(
"Skipping corrupted embedding entry (missing file_path): %s",
embedding_data.get("file_name", "<unknown>"),
)
return None
# Get sub_type from cache entry (new canonical field)
sub_type = embedding_data.get("sub_type", "embedding")
file_name = embedding_data.get("file_name") or ""
model_name = embedding_data.get("model_name") or file_name
folder = embedding_data.get("folder") or ""
return {
"model_name": embedding_data["model_name"],
"file_name": embedding_data["file_name"],
"model_name": model_name,
"file_name": file_name,
"preview_url": config.get_preview_static_url(embedding_data.get("preview_url", "")),
"preview_nsfw_level": embedding_data.get("preview_nsfw_level", 0),
"base_model": embedding_data.get("base_model", ""),
"folder": embedding_data["folder"],
"folder": folder,
"sha256": embedding_data.get("sha256", ""),
"file_path": embedding_data["file_path"].replace(os.sep, "/"),
"file_path": file_path.replace(os.sep, "/"),
"file_size": embedding_data.get("size", 0),
"modified": embedding_data.get("modified", ""),
"tags": embedding_data.get("tags", []),
@@ -48,6 +65,8 @@ class EmbeddingService(BaseModelService):
"skip_metadata_refresh": bool(embedding_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(embedding_data.get("civitai", {}), minimal=True),
"auto_tags": embedding_data.get("auto_tags") or extract_auto_tags(embedding_data),
"version_count": embedding_data.get("version_count"),
"hf_url": embedding_data.get("hf_url", ""),
}
def find_duplicate_hashes(self) -> Dict:
+18
View File
@@ -25,3 +25,21 @@ class ResourceNotFoundError(RuntimeError):
pass
class LLMNotConfiguredError(RuntimeError):
"""Raised when an LLM-dependent operation is attempted but no provider is configured."""
pass
class LLMRateLimitError(RateLimitError):
"""Raised when the LLM provider rejects a request due to rate limiting."""
pass
class LLMResponseError(RuntimeError):
"""Raised when the LLM returns an unparseable or schema-invalid response."""
pass
+729
View File
@@ -0,0 +1,729 @@
"""Centralized LLM API client with BYOK (bring-your-own-key) provider support.
Reads provider configuration from :class:`SettingsManager` and makes
OpenAI-compatible ``/chat/completions`` calls. Supports any provider that
implements the OpenAI Chat Completions API surface area (OpenAI, Ollama,
vLLM, LM Studio, etc.).
"""
from __future__ import annotations
import asyncio
import json
import logging
from typing import Any, Dict, List, Optional
import aiohttp
from .errors import LLMNotConfiguredError, LLMRateLimitError, LLMResponseError
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Model catalog sourced from opencode's maintained model registry.
# maps provider_id -> list of model IDs.
# ---------------------------------------------------------------------------
_MODEL_CATALOG_URL = "https://models.dev/api.json"
# In-memory cache: maps provider slug -> list of model ID strings.
_catalog_cache: Optional[Dict[str, List[str]]] = None
# Per-model max output token limits parsed from the catalog.
# ``{provider_id: {model_id: max_output_tokens}}``.
_model_output_limits: Dict[str, Dict[str, int]] = {}
_CATALOG_TIMEOUT = aiohttp.ClientTimeout(total=30)
async def _load_model_catalog() -> Dict[str, List[str]]:
"""Fetch and parse the model catalog.
Returns ``{provider_id: [model_id, ...]}`` and also populates
:data:`_model_output_limits` with per-model ``limit.output`` values
for use by :func:`_get_model_max_output`.
The JSON at ``_MODEL_CATALOG_URL`` is a dict keyed by provider slug; each
value has a ``models`` sub-dict keyed by model ID. The result is cached
in memory after the first successful fetch.
Subsequent calls return the cached data immediately.
"""
global _catalog_cache, _model_output_limits
if _catalog_cache is not None:
return _catalog_cache
try:
async with aiohttp.ClientSession(timeout=_CATALOG_TIMEOUT) as session:
async with session.get(_MODEL_CATALOG_URL) as resp:
if resp.status != 200:
logger.warning("Model catalog returned HTTP %s", resp.status)
return _catalog_cache or {}
data = await resp.json()
except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError) as exc:
logger.warning("Failed to fetch model catalog: %s", exc)
return _catalog_cache or {}
if not isinstance(data, dict):
logger.warning("Model catalog is not a dict, got %s", type(data).__name__)
return _catalog_cache or {}
result: Dict[str, List[str]] = {}
output_limits: Dict[str, Dict[str, int]] = {}
for provider_id, provider_info in data.items():
if not isinstance(provider_info, dict):
continue
models_dict = provider_info.get("models")
if not isinstance(models_dict, dict):
continue
model_ids: List[str] = []
provider_limits: Dict[str, int] = {}
for mid, model_info in models_dict.items():
if not isinstance(mid, str):
continue
model_ids.append(mid)
if isinstance(model_info, dict):
limit = model_info.get("limit")
if isinstance(limit, dict):
output = limit.get("output")
if isinstance(output, (int, float)) and output > 0:
provider_limits[mid] = int(output)
if model_ids:
result[provider_id] = model_ids
if provider_limits:
output_limits[provider_id] = provider_limits
_catalog_cache = result
_model_output_limits = output_limits
logger.debug(
"Loaded model catalog: %d providers, %d total models "
"(%d providers have output limits)",
len(result),
sum(len(m) for m in result.values()),
len(output_limits),
)
return result
def _get_model_max_output(provider: str, model: str) -> Optional[int]:
"""Return the model's max output token limit from the catalog, or ``None``.
Returns ``None`` when the provider or model is not found in the catalog
(e.g. local Ollama models, custom models, or user-typed model names).
Callers should fall back to a safe default.
"""
return _model_output_limits.get(provider, {}).get(model)
# Short timeout for Ollama's local API
_OLLAMA_API_TIMEOUT = aiohttp.ClientTimeout(total=8)
async def fetch_ollama_models(api_base: str) -> List[str]:
"""Fetch locally available models from a running Ollama instance.
Uses Ollama's OpenAI-compatible ``GET {api_base}/models`` endpoint.
Returns an empty list if Ollama is not reachable (not running).
"""
url = f"{api_base.rstrip('/')}/models"
try:
async with aiohttp.ClientSession(timeout=_OLLAMA_API_TIMEOUT) as session:
async with session.get(url) as resp:
if resp.status != 200:
logger.debug("Ollama API returned HTTP %s from %s", resp.status, api_base)
return []
data = await resp.json()
except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError) as exc:
logger.debug("Ollama not reachable at %s: %s", api_base, exc)
return []
raw = data.get("data") if isinstance(data, dict) else None
if not isinstance(raw, list):
return []
return [
str(entry["id"]) for entry in raw
if isinstance(entry, dict) and isinstance(entry.get("id"), str)
]
async def get_provider_model_ids(provider_id: str) -> List[str]:
"""Return the list of known model IDs for *provider_id* from the catalog.
The catalog is loaded on first call and cached thereafter. If the
provider is not found an empty list is returned (never raises).
"""
catalog = await _load_model_catalog()
return catalog.get(provider_id, [])
async def get_all_provider_models(
provider_ids: List[str],
) -> Dict[str, List[str]]:
"""Return model lists for a subset of providers in one call.
Loads the catalog (cached) and returns only the requested providers.
Handy for embedding lightweight data into the template context.
"""
catalog = await _load_model_catalog()
return {
pid: catalog.get(pid, [])
for pid in provider_ids
}
# Provider preset definitions.
# Each entry contains display metadata and defaults for the UI.
# The key is the internal provider id stored in ``llm_provider``.
# Models are NOT listed here — they come from the opencode model catalog at
# runtime (see :func:`get_provider_model_ids`).
PROVIDER_PRESETS: Dict[str, Dict[str, Any]] = {
"openai": {
"name": "OpenAI",
"api_base": "https://api.openai.com/v1",
"requires_key": True,
},
"ollama": {
"name": "Ollama (local)",
"api_base": "http://localhost:11434/v1",
"requires_key": False,
},
"deepseek": {
"name": "DeepSeek",
"api_base": "https://api.deepseek.com/v1",
"requires_key": True,
},
"groq": {
"name": "Groq",
"api_base": "https://api.groq.com/openai/v1",
"requires_key": True,
},
"openrouter": {
"name": "OpenRouter",
"api_base": "https://openrouter.ai/api/v1",
"requires_key": True,
},
"opencode-go": {
"name": "OpenCode Go",
"api_base": "https://opencode.ai/zen/go/v1",
"requires_key": True,
},
# "custom" is handled specially (no preset api_base, requires user input)
}
# Legacy lookup derived from PROVIDER_PRESETS for backward compat.
_PROVIDER_DEFAULTS: Dict[str, str] = {
pid: info["api_base"]
for pid, info in PROVIDER_PRESETS.items()
if info.get("api_base")
}
# Request timeout for LLM calls (seconds)
_LLM_TIMEOUT = aiohttp.ClientTimeout(total=120)
class LLMService:
"""Centralized LLM API client.
All LLM-based enrichment features call through this service so
that BYOK config, retry logic, and error handling live in one place.
"""
_instance: Optional["LLMService"] = None
_lock: asyncio.Lock = asyncio.Lock()
def __init__(self, settings_service) -> None:
self._settings = settings_service
# ------------------------------------------------------------------
# Singleton access
# ------------------------------------------------------------------
@classmethod
async def get_instance(cls) -> "LLMService":
"""Return the lazily-initialised global ``LLMService`` instance."""
if cls._instance is None:
async with cls._lock:
if cls._instance is None:
from .settings_manager import get_settings_manager
cls._instance = cls(get_settings_manager())
# Start preloading the model catalog in the background so
# the settings UI never blocks on it. The catalog is
# cached after the first fetch (see _load_model_catalog).
asyncio.create_task(_load_model_catalog())
return cls._instance
@classmethod
def reset_instance(cls) -> None:
"""Reset the cached singleton — primarily for tests."""
cls._instance = None
# ------------------------------------------------------------------
# Configuration helpers
# ------------------------------------------------------------------
def _get_config(self) -> Dict[str, Any]:
"""Read the current LLM configuration from settings."""
return {
"provider": self._settings.get("llm_provider", "openai"),
"api_key": self._settings.get("llm_api_key", ""),
"api_base": self._settings.get("llm_api_base", ""),
"model": self._settings.get("llm_model", ""),
}
@staticmethod
def _provider_requires_key(provider: str) -> bool:
"""Return ``False`` when the given provider id does not need an API key."""
preset = PROVIDER_PRESETS.get(provider, {})
return bool(preset.get("requires_key", True))
def is_configured(self) -> bool:
"""Return ``True`` when the LLM provider is minimally configured.
A provider is considered configured when ``llm_model`` is set,
an API key is configured for providers that require one (e.g.
Ollama does not), and an API base URL is set for providers that
have no preset default (e.g. ``custom``).
"""
cfg = self._get_config()
has_model = bool(cfg["model"])
has_key = bool(cfg["api_key"]) or not self._provider_requires_key(cfg["provider"])
has_base = bool(cfg["api_base"]) or bool(_PROVIDER_DEFAULTS.get(cfg["provider"]))
return has_model and has_key and has_base
def _resolve_api_base(self, provider: str, api_base: str) -> str:
"""Resolve the API base URL for the given provider.
If ``api_base`` is explicitly set (non-empty), it takes priority.
Otherwise the default from :data:`PROVIDER_PRESETS` is used.
"""
if api_base:
return api_base.rstrip("/")
return _PROVIDER_DEFAULTS.get(provider, "").rstrip("/")
def _build_headers(self, api_key: str) -> Dict[str, str]:
"""Build HTTP headers for the LLM API request."""
headers = {"Content-Type": "application/json"}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
return headers
def _ensure_configured(self) -> Dict[str, Any]:
"""Validate configuration and return it, or raise.
A provider is considered configured when ``llm_model`` is set,
an API key is configured for providers that require one, and
an API base URL is set for providers without a preset default.
"""
cfg = self._get_config()
has_model = bool(cfg["model"])
needs_key = self._provider_requires_key(cfg["provider"])
has_key = bool(cfg["api_key"]) or not needs_key
has_base = bool(cfg["api_base"]) or bool(_PROVIDER_DEFAULTS.get(cfg["provider"]))
if not (has_model and has_key and has_base):
parts = []
if not has_model:
parts.append("No LLM model specified")
if not has_key and needs_key:
parts.append("No LLM API key configured")
if not has_base:
parts.append(
f"No API base URL for provider '{cfg['provider']}'"
)
detail = "; ".join(parts) if parts else "LLM provider is not configured"
raise LLMNotConfiguredError(
f"{detail}. Configure it in Settings → AI Provider."
)
return cfg
# ------------------------------------------------------------------
# Core API call
# ------------------------------------------------------------------
async def chat_completion(
self,
*,
messages: List[Dict[str, str]],
model: Optional[str] = None,
temperature: float = 0.3,
response_format: Optional[Dict[str, Any]] = None,
max_tokens: Optional[int] = None,
retry_on_rate_limit: bool = True,
) -> Dict[str, Any]:
"""Call the configured LLM provider's ``/chat/completions`` endpoint.
Args:
messages: OpenAI-format message list
model: Override the configured model name
temperature: Sampling temperature
response_format: Optional ``{"type": "json_object"}`` for structured output
max_tokens: Optional max output tokens
retry_on_rate_limit: Retry once after a 429 with backoff
Returns:
Dict with ``content`` (str), ``usage`` (dict), ``model`` (str)
Raises:
LLMNotConfiguredError: Provider not enabled / missing config
LLMRateLimitError: Rate limited and retry exhausted
LLMResponseError: Non-200 response or parse failure
"""
cfg = self._ensure_configured()
api_base = self._resolve_api_base(cfg["provider"], cfg["api_base"])
model_name = model or cfg["model"]
is_ollama = cfg["provider"] == "ollama"
if is_ollama:
# Use Ollama's native /api/chat endpoint which does NOT expose
# a separate reasoning/thinking field (the model's full output
# lands directly in message.content). The OpenAI-compatible
# endpoint splits thinking into the "reasoning" field, making
# content empty when thinking consumes all available tokens.
base = api_base.rstrip("/")
if base.endswith("/v1"):
base = base[:-3]
url = f"{base}/api/chat"
else:
url = f"{api_base}/chat/completions"
payload: Dict[str, Any]
if is_ollama:
payload = {
"model": model_name,
"messages": messages,
"stream": False,
# Suppress separate thinking trace — thinking still happens
# internally (accuracy preserved) but output goes directly to
# message.content instead of being split across content +
# thinking. Without this the model can exhaust num_predict
# on thinking alone and leave content empty.
"think": False,
"options": {
"temperature": temperature,
# 8K context is sufficient for metadata enrichment
# (prompt ~2-5K, output ~0.2-1K tokens). The old 32K
# value was excessive for this use case and increased
# Ollama VRAM usage unnecessarily.
"num_ctx": 8192,
},
}
if response_format is not None:
payload["format"] = "json"
if max_tokens is not None:
payload["options"]["num_predict"] = max_tokens
else:
payload = {
"model": model_name,
"messages": messages,
"temperature": temperature,
}
if response_format is not None:
payload["response_format"] = response_format
if max_tokens is not None:
payload["max_tokens"] = max_tokens
if is_ollama:
logger.info(
"Ollama request: model=%s num_ctx=%s num_predict=%s format=%s think=%s",
payload.get("model"),
payload.get("options", {}).get("num_ctx"),
payload.get("options", {}).get("num_predict"),
payload.get("format", "none"),
payload.get("think"),
)
headers = self._build_headers(cfg["api_key"])
attempt = 0
max_attempts = 2 if retry_on_rate_limit else 1
while attempt < max_attempts:
attempt += 1
try:
async with aiohttp.ClientSession(timeout=_LLM_TIMEOUT) as session:
async with session.post(
url, json=payload, headers=headers
) as resp:
if resp.status == 429:
if attempt < max_attempts:
retry_after = float(
resp.headers.get("Retry-After", "5")
)
logger.warning(
"LLM rate limited, retrying after %.1fs",
retry_after,
)
await asyncio.sleep(retry_after)
continue
raise LLMRateLimitError(
f"LLM provider rate limited (HTTP 429)",
provider=cfg["provider"],
)
if resp.status != 200:
body = await resp.text()
raise LLMResponseError(
f"LLM API returned HTTP {resp.status}: "
f"{body[:500]}"
)
data = await resp.json()
except aiohttp.ClientError as exc:
raise LLMResponseError(f"Network error calling LLM API: {exc}") from exc
# Parse response
try:
if is_ollama:
content = (data.get("message") or {}).get("content") or ""
usage = {"completion_tokens": data.get("eval_count", 0)}
finish_reason = data.get("done_reason", "")
if not content:
logger.warning(
"LLM returned empty content. Provider=ollama, "
"done_reason=%s, eval_count=%s",
finish_reason,
data.get("eval_count", 0),
)
else:
content = data["choices"][0]["message"].get("content") or ""
usage = data.get("usage", {})
if not content:
logger.warning(
"LLM returned empty content. Full response truncated: %s",
json.dumps(data, ensure_ascii=False)[:1000],
)
return {
"content": content,
"usage": usage,
"model": data.get("model", model_name),
}
except (KeyError, IndexError) as exc:
raise LLMResponseError(
f"Unexpected LLM response structure: {json.dumps(data)[:500]}"
) from exc
# Should not reach here, but satisfy type checker
raise LLMRateLimitError("Rate limit retry exhausted", provider=cfg["provider"])
# ------------------------------------------------------------------
# Structured output convenience
# ------------------------------------------------------------------
async def chat_completion_json(
self,
*,
system_prompt: str,
user_prompt: str,
model: Optional[str] = None,
temperature: float = 0.3,
max_tokens: Optional[int] = None,
) -> Dict[str, Any]:
"""Call the LLM with ``response_format=json_object`` and return parsed JSON.
``max_tokens`` is resolved in this order:
1. Explicit caller-supplied ``max_tokens``
2. Per-model ``limit.output`` from the model catalog
3. A safe default of 4096 (sufficient for metadata enrichment)
If the response content is empty or not valid JSON, attempts
:func:`_try_salvage_json` before raising.
Args:
system_prompt: System-level instructions
user_prompt: User-level query
model: Override the configured model name
temperature: Sampling temperature
max_tokens: Optional max output tokens
Returns:
Parsed JSON dict from the LLM response
Raises:
LLMNotConfiguredError: Provider not configured
LLMRateLimitError: Rate limited
LLMResponseError: Empty response or JSON parse failure
"""
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
]
# Resolve max_tokens: caller override → catalog lookup → safe default
if max_tokens is None:
cfg = self._get_config()
effective_max = _get_model_max_output(cfg["provider"], cfg["model"])
else:
effective_max = max_tokens
if effective_max is None:
effective_max = 4096
# Use json_schema (not json_object) for broader provider compatibility:
# LM Studio and some other OpenAI-compatible servers reject
# json_object but accept json_schema. {"type": "object"} is
# functionally equivalent — it accepts any JSON object without
# constraining specific fields.
response_format = {
"type": "json_schema",
"json_schema": {
"name": "metadata",
"schema": {"type": "object"},
},
}
try:
result = await self.chat_completion(
messages=messages,
model=model,
temperature=temperature,
response_format=response_format,
max_tokens=effective_max,
)
except LLMResponseError as e:
# Only fall back when the provider rejects the response_format
# type value (e.g. "'response_format.type' must be..."). Avoid
# catching unrelated 400 errors whose body happens to mention
# "response_format" (e.g. "model does not support
# response_format restrictions on this endpoint").
if "'response_format.type'" not in str(e).lower():
raise
logger.info(
"Provider rejected response_format, retrying without it. "
"Falling back to prompt-only JSON mode. Error: %s",
e,
)
result = await self.chat_completion(
messages=messages,
model=model,
temperature=temperature,
response_format=None,
max_tokens=effective_max,
)
content = result.get("content", "") or ""
if not content:
raise LLMResponseError(
"LLM returned empty content. "
f"Raw response: {json.dumps(result)[:500]}"
)
try:
parsed = json.loads(content)
logger.debug(
"LLM raw content: %s",
json.dumps(parsed, ensure_ascii=False)[:2000],
)
return parsed
except (json.JSONDecodeError, TypeError) as exc:
logger.info(
"LLM raw response (first 800 chars): %s",
content[:800],
)
# Last resort: attempt to salvage partial/truncated JSON
salvaged = _try_salvage_json(content)
if salvaged is not None:
logger.warning(
"LLM JSON salvaged from partial content (%d chars raw)",
len(content),
)
return salvaged
raise LLMResponseError(
f"LLM response could not be parsed as JSON: {content[:200]}"
)
def _try_salvage_json(raw: str) -> Dict[str, Any] | None:
"""Attempt to repair and parse a truncated JSON string.
Handles common truncation patterns:
* Incomplete string value at the end (``"foo`` → ``"foo"``)
* Missing closing ``}`` or ``]`` (respecting nesting order)
* Trailing comma before closing bracket
* Extra text after the JSON object (e.g. markdown fences)
Returns the parsed dict on success, ``None`` if repair is impossible.
"""
if not raw:
return None
text = raw.strip()
# Strip markdown fences if the LLM wrapped the JSON
if text.startswith("```"):
end = text.find("\n")
text = text[end + 1:] if end != -1 else text[3:]
if text.endswith("```"):
text = text[:-3].rstrip()
# Find the first '{' and strip everything before it
start = text.find("{")
if start == -1:
return None
text = text[start:]
# Try to close an incomplete string at the end (e.g. ``"https://huggingf``)
# Pattern: ends mid-string (last quote is open)
if text.count('"') % 2 == 1:
text += '"'
# Ensure trailing commas before closing braces work
text = _strip_trailing_commas(text)
# Walk through the text character by character to find unclosed
# brackets and close them in the correct (LIFO) order.
# We ignore brackets inside quoted strings.
stack: list[str] = []
in_string = False
escape = False
for ch in text:
if escape:
escape = False
continue
if ch == "\\":
escape = True
continue
if ch == '"':
in_string = not in_string
continue
if in_string:
continue
if ch in ("{", "["):
stack.append(ch)
elif ch == "}":
if stack and stack[-1] == "{":
stack.pop()
else:
return None # Unmatched closer — unrecoverable
elif ch == "]":
if stack and stack[-1] == "[":
stack.pop()
else:
return None
# Close remaining open brackets in reverse order
for opener in reversed(stack):
text += "}" if opener == "{" else "]"
try:
return json.loads(text)
except (json.JSONDecodeError, ValueError):
return None
def _strip_trailing_commas(text: str) -> str:
"""Remove commas that appear before a closing brace/bracket."""
import re as _re
text = _re.sub(r",\s*}", "}", text)
text = _re.sub(r",\s*]", "]", text)
return text
+33 -9
View File
@@ -24,23 +24,41 @@ class LoraService(BaseModelService):
"""
super().__init__("lora", scanner, LoraMetadata, update_service=update_service)
async def format_response(self, lora_data: Dict) -> Dict:
"""Format LoRA data for API response"""
async def format_response(self, lora_data: Dict) -> Optional[Dict]:
"""Format LoRA data for API response.
Returns None when the entry is missing critical fields (corrupted cache
row), so the handler layer can filter it out instead of crashing the
whole listing request. See issue #730.
"""
# Guard against corrupted cache entries missing critical fields
file_path = lora_data.get("file_path")
if not file_path or not isinstance(file_path, str):
logger.warning(
"Skipping corrupted LoRA entry (missing file_path): %s",
lora_data.get("file_name", "<unknown>"),
)
return None
# Resolve sub_type using priority: sub_type > model_type > civitai.model.type > default
# Normalize to lowercase for consistent API responses
sub_type = resolve_sub_type(lora_data).lower()
file_name = lora_data.get("file_name") or ""
model_name = lora_data.get("model_name") or file_name
folder = lora_data.get("folder") or ""
return {
"model_name": lora_data["model_name"],
"file_name": lora_data["file_name"],
"model_name": model_name,
"file_name": file_name,
"preview_url": config.get_preview_static_url(
lora_data.get("preview_url", "")
),
"preview_nsfw_level": lora_data.get("preview_nsfw_level", 0),
"base_model": lora_data.get("base_model", ""),
"folder": lora_data["folder"],
"folder": folder,
"sha256": lora_data.get("sha256", ""),
"file_path": lora_data["file_path"].replace(os.sep, "/"),
"file_path": file_path.replace(os.sep, "/"),
"file_size": lora_data.get("size", 0),
"modified": lora_data.get("modified", ""),
"tags": lora_data.get("tags", []),
@@ -59,6 +77,8 @@ class LoraService(BaseModelService):
lora_data.get("civitai", {}), minimal=True
),
"auto_tags": lora_data.get("auto_tags") or extract_auto_tags(lora_data),
"version_count": lora_data.get("version_count"),
"hf_url": lora_data.get("hf_url", ""),
}
async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]:
@@ -251,12 +271,16 @@ class LoraService(BaseModelService):
return letters
async def get_lora_trigger_words(self, lora_name: str) -> List[str]:
"""Get trigger words for a specific LoRA file"""
"""Get trigger words for a specific LoRA file.
Supports both simple names and full-path syntax.
"""
cache = await self.scanner.get_cached_data()
for lora in cache.raw_data:
if lora["file_name"] == lora_name:
civitai_data = lora.get("civitai", {})
file_name = lora.get("file_name", "")
if file_name == lora_name or lora_name.endswith("/" + file_name) or lora_name.endswith("\\" + file_name):
civitai_data = lora.get("civitai") or {}
return civitai_data.get("trainedWords", [])
return []
+69 -16
View File
@@ -15,6 +15,17 @@ from .service_registry import ServiceRegistry
logger = logging.getLogger(__name__)
_PROVIDER_DISPLAY_NAMES = {
"civitai_api": "CivitAI",
"civarchive_api": "CivArchive",
"sqlite": "Archive DB",
}
_PRESET_PROVIDER_ORDERS = {
"civitai_archive_sqlite": ["civitai_api", "civarchive_api", "sqlite"],
"civitai_sqlite_archive": ["civitai_api", "sqlite", "civarchive_api"],
}
async def initialize_metadata_providers():
"""Initialize and configure all metadata providers based on settings"""
provider_manager = await ModelMetadataProviderManager.get_instance()
@@ -26,7 +37,9 @@ async def initialize_metadata_providers():
# Get settings
settings_manager = get_settings_manager()
enable_archive_db = settings_manager.get('enable_metadata_archive_db', False)
enable_civarchive_api = settings_manager.get('enable_civarchive_api', True)
provider_order = settings_manager.get('metadata_provider_order', 'civitai_archive_sqlite')
providers = []
# Initialize archive database provider if enabled
@@ -59,27 +72,48 @@ async def initialize_metadata_providers():
except Exception as e:
logger.error(f"Failed to initialize Civitai API metadata provider: {e}")
# Register CivArchive provider, and all add to fallback providers
try:
civarchive_client = await ServiceRegistry.get_civarchive_client()
civarchive_provider = CivArchiveModelMetadataProvider(civarchive_client)
provider_manager.register_provider('civarchive_api', civarchive_provider)
providers.append(('civarchive_api', civarchive_provider))
logger.debug("CivArchive metadata provider registered (also included in fallback)")
except Exception as e:
logger.error(f"Failed to initialize CivArchive metadata provider: {e}")
# Register CivArchive provider when enabled. Civitai API is always
# preferred (better metadata); CivArchive mainly recovers metadata for
# models deleted from Civitai, so it can be turned off to avoid its long
# rate-limit windows entirely.
if enable_civarchive_api:
try:
civarchive_client = await ServiceRegistry.get_civarchive_client()
civarchive_provider = CivArchiveModelMetadataProvider(civarchive_client)
provider_manager.register_provider('civarchive_api', civarchive_provider)
providers.append(('civarchive_api', civarchive_provider))
logger.debug("CivArchive metadata provider registered (also included in fallback)")
except Exception as e:
logger.error(f"Failed to initialize CivArchive metadata provider: {e}")
else:
logger.debug("CivArchive metadata provider disabled by setting 'enable_civarchive_api'")
# Preset fallback orderings (see module-level _PRESET_PROVIDER_ORDERS).
# civitai_api is always first (better metadata); the remaining providers
# are arranged by the configured preset. Providers that are not
# registered (disabled/unavailable) are simply skipped, so each preset
# degrades gracefully.
desired_order = _PRESET_PROVIDER_ORDERS.get(
provider_order, _PRESET_PROVIDER_ORDERS["civitai_archive_sqlite"]
)
# Set up fallback provider based on available providers
if len(providers) > 1:
# Always use Civitai API (it has better metadata), then CivArchive API, then Archive DB
ordered_providers: list[tuple[str, ModelMetadataProvider]] = []
ordered_providers.extend([p for p in providers if p[0] == 'civitai_api'])
ordered_providers.extend([p for p in providers if p[0] == 'civarchive_api'])
ordered_providers.extend([p for p in providers if p[0] == 'sqlite'])
for name in desired_order:
ordered_providers.extend([p for p in providers if p[0] == name])
# Include any provider not covered by the preset (defensive) at the end
for p in providers:
if p not in ordered_providers:
ordered_providers.append(p)
if ordered_providers:
fallback_provider = FallbackMetadataProvider(ordered_providers)
provider_manager.register_provider('fallback', fallback_provider, is_default=True)
logger.debug(
"Metadata fallback provider order: %s",
", ".join(name for name, _ in ordered_providers),
)
elif len(providers) == 1:
# Only one provider available, set it as default
provider_name, provider = providers[0]
@@ -96,11 +130,30 @@ async def update_metadata_providers():
# Get current settings
settings_manager = get_settings_manager()
enable_archive_db = settings_manager.get('enable_metadata_archive_db', False)
enable_civarchive_api = settings_manager.get('enable_civarchive_api', True)
provider_order = settings_manager.get('metadata_provider_order', 'civitai_archive_sqlite')
# Reinitialize all providers with new settings
provider_manager = await initialize_metadata_providers()
logger.info(f"Updated metadata providers, archive_db enabled: {enable_archive_db}")
# Build effective provider chain for logging (use actually-registered
# providers, not just settings, so a failed init is reflected correctly)
registered = set(provider_manager.providers.keys())
desired = _PRESET_PROVIDER_ORDERS.get(
provider_order, _PRESET_PROVIDER_ORDERS["civitai_archive_sqlite"]
)
chain = "".join(
_PROVIDER_DISPLAY_NAMES[p]
for p in desired
if p in registered and p in _PROVIDER_DISPLAY_NAMES
)
logger.info(
"Updated metadata providers: archive_db=%s, civarchive_api=%s, chain=%s",
enable_archive_db,
enable_civarchive_api,
chain,
)
return provider_manager
except Exception as e:
logger.error(f"Failed to update metadata providers: {e}")
+50 -9
View File
@@ -209,20 +209,40 @@ class MetadataSyncService:
error_msg = "CivitAI model is deleted and no archive provider is available"
return False, error_msg
else:
provider_attempts.append((None, await self._get_default_provider()))
is_hf_source = bool(model_data.get("hf_url"))
if is_hf_source:
# HF-sourced model: only check CivitAI API directly.
# CivArchive is almost guaranteed to have no record, and
# hitting it wastes rate-limit budget.
# Use a distinct provider name ("civitai_api" not None) so
# downstream code does NOT interpret a "Model not found"
# response as civitai_api_not_found — which would mark the
# model civitai_deleted=True when it was never on CivitAI.
try:
provider_attempts.append(("civitai_api", await self._get_provider("civitai_api")))
except Exception as exc: # pragma: no cover - provider resolution fault
logger.debug("Unable to resolve civitai_api provider: %s", exc)
if not provider_attempts:
provider_attempts.append((None, await self._get_default_provider()))
civitai_metadata: Optional[Dict[str, Any]] = None
metadata_provider: Optional[MetadataProviderProtocol] = None
provider_used: Optional[str] = None
last_error: Optional[str] = None
civitai_api_not_found = False
any_rate_limited = False
for provider_name, provider in provider_attempts:
try:
civitai_metadata_candidate, error = await provider.get_model_by_hash(sha256)
except RateLimitError as exc:
exc.provider = exc.provider or (provider_name or provider.__class__.__name__)
raise
logger.warning(
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
provider_name or provider.__class__.__name__,
exc.retry_after or 0,
)
any_rate_limited = True
continue
except Exception as exc: # pragma: no cover - defensive logging
logger.error("Provider %s failed for hash %s: %s", provider_name, sha256, exc)
civitai_metadata_candidate, error = None, str(exc)
@@ -258,6 +278,14 @@ class MetadataSyncService:
model_data["last_checked_at"] = datetime.now().timestamp()
needs_save = True
# When the model was already classified as "not on CivitAI" via
# .metadata.json (civitai_deleted=True) but the SQLite cache is
# stale (because the pre-fix code never persisted these flags),
# ensure the flags are written to the scanner cache + SQLite.
if not needs_save and model_data.get("civitai_deleted") is True:
model_data["last_checked_at"] = datetime.now().timestamp()
needs_save = True
# Save metadata if any state was updated
if needs_save:
data_to_save = model_data.copy()
@@ -266,6 +294,7 @@ class MetadataSyncService:
if "last_checked_at" not in data_to_save:
data_to_save["last_checked_at"] = datetime.now().timestamp()
await self._metadata_manager.save_metadata(file_path, data_to_save)
await update_cache_func(file_path, file_path, data_to_save)
default_error = (
"CivitAI model is deleted and metadata archive DB is not enabled"
@@ -276,17 +305,18 @@ class MetadataSyncService:
)
resolved_error = last_error or default_error
if any_rate_limited and "Rate limited" not in resolved_error:
resolved_error = "Rate limited"
if is_expected_offline_error(resolved_error):
resolved_error = OFFLINE_FRIENDLY_MESSAGE
error_msg = (
f"Error fetching metadata: {resolved_error} "
f"(model_name={model_data.get('model_name', '')})"
f"(file={os.path.basename(file_path)}, sha256={sha256})"
)
if is_expected_offline_error(resolved_error):
logger.info(error_msg)
else:
logger.error(error_msg)
# Use case layer (BulkMetadataRefreshUseCase) logs failed models at WARNING level,
# so this level is demoted to DEBUG to avoid duplicate user-visible logging.
logger.debug(error_msg)
return False, error_msg
model_data["from_civitai"] = True
@@ -411,7 +441,18 @@ class MetadataSyncService:
metadata = await metadata_loader(metadata_path)
for key, value in updates.items():
if isinstance(value, dict) and isinstance(metadata.get(key), dict):
if key == "tags" and isinstance(value, list):
# Normalize tags: trim, lowercase, deduplicate
normalized = []
seen = set()
for tag in value:
if isinstance(tag, str):
t = tag.strip().lower()
if t and t not in seen:
normalized.append(t)
seen.add(t)
metadata[key] = normalized
elif isinstance(value, dict) and isinstance(metadata.get(key), dict):
metadata[key].update(value)
else:
metadata[key] = value
+35 -1
View File
@@ -18,6 +18,8 @@ SUPPORTED_SORT_MODES = [
('size', 'desc'),
('usage', 'asc'),
('usage', 'desc'),
('versions_count', 'asc'),
('versions_count', 'desc'),
]
# Is this in use?
@@ -263,6 +265,17 @@ class ModelCache:
),
reverse=reverse
)
elif sort_key == 'versions_count':
# Pre-dedup sort: fall back to name sort.
# Actual re-sort by version_count happens in get_paginated_data after dedup.
result = natsorted(
data,
key=lambda x: (
self._get_display_name(x).lower(),
x.get('file_path', '').lower()
),
reverse=reverse
)
else:
# Fallback: no sort
result = list(data)
@@ -324,4 +337,25 @@ class ModelCache:
else:
return False # Model not found
return True
return True
async def clear_preview_by_path(self, preview_file_path: str) -> int:
"""Clear ``preview_url`` for every cached entry referencing a file path.
When a preview file has been deleted from disk, this removes its
reference from all matching cache entries so the next list-API
response returns an empty ``preview_url`` instead of a stale URL
that produces 404s.
Returns the number of entries that were updated.
"""
normalized = preview_file_path.replace("\\", "/")
cleared = 0
async with self._lock:
for item in self.raw_data:
cached_url = item.get("preview_url", "")
if cached_url.replace("\\", "/") == normalized:
item["preview_url"] = ""
item["preview_nsfw_level"] = 0
cleared += 1
return cleared
+4
View File
@@ -8,6 +8,7 @@ from abc import ABC, abstractmethod
from ..utils.utils import calculate_relative_path_for_model, remove_empty_dirs
from ..utils.constants import AUTO_ORGANIZE_BATCH_SIZE
from ..services.settings_manager import get_settings_manager
from ..services.model_lifecycle_service import _require_path_in_library_roots
logger = logging.getLogger(__name__)
@@ -493,6 +494,9 @@ class ModelMoveService:
Dictionary with move result
"""
try:
_require_path_in_library_roots(file_path, self.scanner, label="Source path")
_require_path_in_library_roots(target_path, self.scanner, label="Target path")
if use_default_paths:
# Find the model in cache to get metadata
cache = await self.scanner.get_cached_data()
+41
View File
@@ -48,6 +48,36 @@ async def delete_model_artifacts(
return deleted
def _require_path_in_library_roots(file_path: str, scanner, *, label: str = "path") -> None:
"""Raise ``ValueError`` if *file_path* is not inside a configured model root.
Uses ``os.path.abspath()`` (NOT ``realpath``) to resolve ``..`` and ``.``
while preserving symlinks this keeps the check in business-path space.
Skips when the scanner does not expose ``get_model_roots`` or the list
is empty.
"""
roots = None
if hasattr(scanner, "get_model_roots"):
try:
roots = scanner.get_model_roots()
except NotImplementedError:
roots = None
if not roots:
return
resolved = os.path.abspath(os.path.normpath(file_path))
for root in roots:
root_resolved = os.path.abspath(os.path.normpath(root))
if resolved == root_resolved or resolved.startswith(root_resolved + os.sep):
return
raise ValueError(
f"{label} '{file_path}' is outside configured library directories"
)
class ModelLifecycleService:
"""Co-ordinate destructive and mutating model operations."""
@@ -74,6 +104,8 @@ class ModelLifecycleService:
if not file_path:
raise ValueError("Model path is required")
_require_path_in_library_roots(file_path, self._scanner, label="File path")
cache = await self._scanner.get_cached_data()
cached_entry = None
@@ -182,6 +214,8 @@ class ModelLifecycleService:
if not file_path:
raise ValueError("Model path is required")
_require_path_in_library_roots(file_path, self._scanner, label="File path")
metadata_path = os.path.splitext(file_path)[0] + ".metadata.json"
metadata = await self._metadata_loader(metadata_path)
metadata["exclude"] = True
@@ -229,6 +263,8 @@ class ModelLifecycleService:
if not file_path:
raise ValueError("Model path is required")
_require_path_in_library_roots(file_path, self._scanner, label="File path")
if not os.path.exists(file_path):
raise ValueError("Model file does not exist")
@@ -270,6 +306,9 @@ class ModelLifecycleService:
if not file_paths:
raise ValueError("No file paths provided for deletion")
for path in file_paths:
_require_path_in_library_roots(path, self._scanner, label="File path")
return await self._scanner.bulk_delete_models(file_paths)
async def rename_model(
@@ -280,6 +319,8 @@ class ModelLifecycleService:
if not file_path or not new_file_name:
raise ValueError("File path and new file name are required")
_require_path_in_library_roots(file_path, self._scanner, label="File path")
invalid_chars = {"/", "\\", ":", "*", "?", '"', "<", ">", "|"}
if any(char in new_file_name for char in invalid_chars):
raise ValueError("Invalid characters in file name")
+44 -21
View File
@@ -65,7 +65,14 @@ class _RateLimitRetryHelper:
return await func(*args, **kwargs)
except RateLimitError as exc:
attempt += 1
if attempt >= self._retry_limit:
# Determine effective retry limit based on rate-limit magnitude
effective_retry_limit = self._retry_limit # default: 3
if exc.retry_after is not None and exc.retry_after >= 120.0:
# Long rate-limit window (>=2 min) — retries are futile
effective_retry_limit = 1 # total 1 attempt = 0 retries
if attempt >= effective_retry_limit:
exc.provider = exc.provider or label
raise
@@ -478,8 +485,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result:
return result, error
except RateLimitError as exc:
exc.provider = exc.provider or label
raise exc
logger.warning(
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e:
logger.debug("Provider %s failed for get_model_by_hash: %s", label, e)
continue
@@ -497,16 +508,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result:
return result
except RateLimitError as exc:
if not_found_confirmed:
logger.debug(
"Suppressing rate limit from %s for model %s: "
"already confirmed as not found by another provider",
label,
model_id,
)
return None
exc.provider = exc.provider or label
raise exc
logger.warning(
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except ResourceNotFoundError:
not_found_confirmed = True
logger.debug(
@@ -532,8 +539,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result:
return result
except RateLimitError as exc:
exc.provider = exc.provider or label
raise exc
logger.warning(
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e:
logger.debug("Provider %s failed for get_model_version: %s", label, e)
continue
@@ -550,8 +561,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result:
return result, error
except RateLimitError as exc:
exc.provider = exc.provider or label
raise exc
logger.warning(
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e:
logger.debug("Provider %s failed for get_model_version_info: %s", label, e)
continue
@@ -572,8 +587,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
except NotImplementedError:
continue
except RateLimitError as exc:
exc.provider = exc.provider or label
raise exc
logger.warning(
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e:
logger.debug(
"Provider %s failed for get_model_versions_by_hashes: %s",
@@ -594,8 +613,12 @@ class FallbackMetadataProvider(ModelMetadataProvider):
if result is not None:
return result
except RateLimitError as exc:
exc.provider = exc.provider or label
raise exc
logger.warning(
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
label,
exc.retry_after or 0,
)
continue
except Exception as e:
logger.debug("Provider %s failed for get_user_models: %s", label, e)
continue
+17 -9
View File
@@ -294,12 +294,14 @@ class ModelFilterSet:
for tag, state in tag_filters.items():
if not tag:
continue
# Normalize to lowercase for case-insensitive matching
normalized = tag.strip().lower()
if state == "exclude":
exclude_tags.add(tag)
exclude_tags.add(normalized)
else:
include_tags.add(tag)
include_tags.add(normalized)
else:
include_tags = {tag for tag in tag_filters if tag}
include_tags = {tag.strip().lower() for tag in tag_filters if tag}
if include_tags:
tag_logic = criteria.tag_logic.lower() if criteria.tag_logic else "any"
@@ -318,13 +320,17 @@ class ModelFilterSet:
return True
# Otherwise, check if all non-special tags match
if non_special_tags:
return all(tag in (item_tags or []) for tag in non_special_tags)
# Case-insensitive: normalize item tags too
normalized_item_tags = {t.strip().lower() for t in (item_tags or []) if isinstance(t, str)}
return all(tag in normalized_item_tags for tag in non_special_tags)
return True
# Normal case: all tags must match
return all(tag in (item_tags or []) for tag in non_special_tags)
# Normal case: all tags must match (case-insensitive)
normalized_item_tags = {t.strip().lower() for t in (item_tags or []) if isinstance(t, str)}
return all(tag in normalized_item_tags for tag in non_special_tags)
else:
# OR logic (default): item must have ANY include tag
return any(tag in include_tags for tag in (item_tags or []))
# OR logic (default): item must have ANY include tag (case-insensitive)
normalized_item_tags = {t.strip().lower() for t in (item_tags or []) if isinstance(t, str)}
return bool(normalized_item_tags & include_tags)
items = [item for item in items if matches_include(item.get("tags"))]
@@ -333,7 +339,9 @@ class ModelFilterSet:
def matches_exclude(item_tags):
if not item_tags and "__no_tags__" in exclude_tags:
return True
return any(tag in exclude_tags for tag in (item_tags or []))
# Case-insensitive: normalize item tags
normalized_item_tags = {t.strip().lower() for t in (item_tags or []) if isinstance(t, str)}
return bool(normalized_item_tags & exclude_tags)
items = [
item for item in items if not matches_exclude(item.get("tags"))
+334 -21
View File
@@ -14,7 +14,7 @@ from ..utils.metadata_manager import MetadataManager
from ..utils.civitai_utils import resolve_license_info
from .model_cache import ModelCache
from .model_hash_index import ModelHashIndex
from .model_lifecycle_service import delete_model_artifacts
from .model_lifecycle_service import delete_model_artifacts, _require_path_in_library_roots
from .service_registry import ServiceRegistry
from .websocket_manager import ws_manager
from .persistent_model_cache import get_persistent_cache
@@ -227,6 +227,11 @@ class ModelScanner:
entry: Dict[str, Any] = {
'file_path': normalized_path,
# file_name is always stored WITHOUT extension (e.g. "OWSMianne_ANIMA_V1",
# not "OWSMianne_ANIMA_V1.safetensors"). All upstream population points
# (MetadataManager, from_civitai_info, download manager, etc.) strip the
# extension via os.path.splitext before writing. Code consuming this field
# should match against names that are likewise extension-free.
'file_name': get_value('file_name', '') or '',
'model_name': get_value('model_name', '') or '',
'folder': normalized_folder,
@@ -248,6 +253,7 @@ class ModelScanner:
'civitai': civitai_slim,
'civitai_deleted': bool(get_value('civitai_deleted', False)),
'skip_metadata_refresh': bool(get_value('skip_metadata_refresh', False)),
'hf_url': get_value('hf_url', '') or '',
}
license_source: Dict[str, Any] = {}
@@ -476,11 +482,20 @@ class ModelScanner:
for tag in adjusted_item.get('tags') or []:
tags_count[tag] = tags_count.get(tag, 0) + 1
# Validate cache entries and check health
# Validate cache entries and check health.
# Always use the validated/repaired entries — even when there are no
# invalid entries, auto_repair may have filled in missing optional
# fields (model_name, file_name, folder) with safe defaults on a copied
# working_entry. Without this unconditional replacement the repaired
# copies are discarded and None values propagate to format_response.
# See issue #730.
valid_entries, invalid_entries = CacheEntryValidator.validate_batch(
adjusted_raw_data, auto_repair=True
)
# Always use the validated entries (repaired copies)
adjusted_raw_data = valid_entries
if invalid_entries:
monitor = CacheHealthMonitor()
report = monitor.check_health(adjusted_raw_data, auto_repair=True)
@@ -532,6 +547,13 @@ class ModelScanner:
if not scan_result or not getattr(self, '_persistent_cache', None):
return
if self.is_cancelled():
logger.info(
f"{self.model_type.capitalize()} Scanner: Skipping _save_persistent_cache "
"after cancellation"
)
return
hash_snapshot = self._build_hash_index_snapshot(scan_result.hash_index)
loop = asyncio.get_event_loop()
try:
@@ -705,14 +727,20 @@ class ModelScanner:
# Determine the page type based on model type
# Scan for new data
scan_result = await self._gather_model_data()
await self._apply_scan_result(scan_result)
await self._save_persistent_cache(scan_result)
await self._sync_download_history(scan_result.raw_data, source='scan')
if not self.is_cancelled():
await self._apply_scan_result(scan_result)
await self._save_persistent_cache(scan_result)
await self._sync_download_history(scan_result.raw_data, source='scan')
logger.info(
f"{self.model_type.capitalize()} Scanner: Cache initialization completed in {time.time() - start_time:.2f} seconds, "
f"found {len(scan_result.raw_data)} models"
)
logger.info(
f"{self.model_type.capitalize()} Scanner: Cache initialization completed in {time.time() - start_time:.2f} seconds, "
f"found {len(scan_result.raw_data)} models"
)
else:
logger.info(
f"{self.model_type.capitalize()} Scanner: Cache initialization cancelled "
f"after {time.time() - start_time:.2f} seconds"
)
except Exception as e:
logger.error(f"{self.model_type.capitalize()} Scanner: Error initializing cache: {e}")
# Ensure cache is at least an empty structure on error
@@ -899,6 +927,25 @@ class ModelScanner:
# Update cache data
self._cache.raw_data = [item for item in self._cache.raw_data if item['file_path'] not in missing_files]
dedup_removed = 0
seen_paths: set = set()
deduped: list = []
for item in reversed(self._cache.raw_data):
path = item.get('file_path', '')
if path not in seen_paths:
seen_paths.add(path)
deduped.append(item)
else:
for tag in item.get('tags', []):
if tag in self._tags_count:
self._tags_count[tag] = max(0, self._tags_count[tag] - 1)
if self._tags_count[tag] == 0:
del self._tags_count[tag]
dedup_removed += 1
if dedup_removed > 0:
self._cache.raw_data = list(reversed(deduped))
total_removed += dedup_removed
# Resort cache if changes were made
if total_added > 0 or total_removed > 0:
# Update folders list
@@ -1067,8 +1114,11 @@ class ModelScanner:
model_data = self._build_cache_entry(metadata, folder=normalized_folder)
# Compute SHA256 hash when metadata provided none (e.g., CivitAI API response has empty hashes)
if not model_data.get('sha256') and file_path:
# Compute SHA256 hash when metadata provided none (e.g., CivitAI API response has empty hashes).
# Respect hash_status='pending' (set by CheckpointScanner for large models) to defer
# hash calculation until on-demand — avoids reading entire checkpoint files at startup.
hash_status = model_data.get('hash_status', '')
if not model_data.get('sha256') and hash_status != 'pending' and file_path:
try:
logger.info(f"Computing SHA256 hash for {file_path} (was empty from metadata)")
sha256 = await calculate_sha256(file_path)
@@ -1093,6 +1143,13 @@ class ModelScanner:
if scan_result is None:
return
if self.is_cancelled():
logger.info(
f"{self.model_type.capitalize()} Scanner: Skipping _apply_scan_result "
"after cancellation"
)
return
self._hash_index = scan_result.hash_index
self._tags_count = dict(scan_result.tags_count)
self._excluded_models = list(scan_result.excluded_models)
@@ -1314,18 +1371,25 @@ class ModelScanner:
# Update folder in metadata
metadata_dict['folder'] = folder
# Add to cache
self._cache.raw_data.append(metadata_dict)
self._cache.add_to_version_index(metadata_dict)
file_path = metadata_dict.get('file_path', '')
if file_path:
old_entries = [item for item in self._cache.raw_data if item.get('file_path') == file_path]
for old_entry in old_entries:
for tag in old_entry.get('tags', []):
if tag in self._tags_count:
self._tags_count[tag] = max(0, self._tags_count[tag] - 1)
if self._tags_count[tag] == 0:
del self._tags_count[tag]
self._hash_index.remove_by_path(file_path)
self._cache.raw_data = [item for item in self._cache.raw_data if item.get('file_path') != file_path]
for tag in metadata_dict.get('tags', []):
self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
self._cache.raw_data.append(metadata_dict)
# Resort cache data
await self._cache.resort()
# Update folders list
all_folders = set(self._cache.folders)
all_folders.add(folder)
self._cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
# Update the hash index
self._hash_index.add_entry(metadata_dict['sha256'], metadata_dict['file_path'])
await self._persist_current_cache()
@@ -1356,6 +1420,9 @@ class ModelScanner:
base_name = os.path.splitext(os.path.basename(source_path))[0]
source_dir = os.path.dirname(source_path)
_require_path_in_library_roots(source_path, self, label="Source path")
_require_path_in_library_roots(target_path, self, label="Target path")
os.makedirs(target_path, exist_ok=True)
@@ -1528,6 +1595,218 @@ class ModelScanner:
return cache_entry if metadata else True
async def sync_cache_from_metadata(
self, file_path: str, metadata_dict: Dict[str, Any]
) -> bool:
"""Opportunistically sync in-memory and persistent caches from metadata.
Builds a prospective cache entry from *metadata_dict* (deserialized
``.metadata.json`` content) and compares it against the current cache
entry. When the two are already identical this method returns
``False`` without touching anything avoiding the overhead of
``update_single_model_cache``, which always removes and re-inserts
the entry, triggers a full resort, and persists via the heavyweight
``save_cache()``.
When differences are detected the update is applied **in-place** with
targeted operations:
* The existing ``raw_data`` entry is modified rather than removed and
re-appended (O(1) instead of O(n)).
* Tag counts and the hash index are updated incrementally.
* The version index is rebuilt only for the affected entry.
* ``resort()`` is called **only** when a sort-relevant field changed
(``model_name`` / ``file_name`` for name-sort, ``modified`` for
date-sort, ``size`` for size-sort).
* The persistent (SQLite) cache receives a targeted single-row update
via :meth:`PersistentModelCache.update_single_model` rather than a
full-table ``save_cache()``.
Returns:
``True`` if any cache update was performed, ``False`` if the
caches were already in sync.
.. note::
This is a **best-effort** operation. Failures are logged but
never propagated callers should fire-and-forget via
:func:`asyncio.create_task`.
"""
try:
return await self._sync_cache_from_metadata_impl(
file_path, metadata_dict
)
except Exception:
logger.warning(
"sync_cache_from_metadata failed for %s",
file_path,
exc_info=True,
)
return False
async def _sync_cache_from_metadata_impl(
self, file_path: str, metadata_dict: Dict[str, Any]
) -> bool:
cache = await self.get_cached_data()
# Locate the existing cache entry -----------------------------------
existing_idx: Optional[int] = None
existing_entry: Optional[Dict[str, Any]] = None
for i, item in enumerate(cache.raw_data):
if item.get("file_path") == file_path:
existing_entry = item
existing_idx = i
break
# Build the desired entry from metadata ------------------------------
folder_value = (
existing_entry.get("folder", "")
if existing_entry
else self._calculate_folder(file_path)
)
desired_entry = self._build_cache_entry(
metadata_dict,
folder=folder_value,
file_path_override=file_path,
)
# Ensure sha256 is populated (defensive — metadata should have it)
if (
not desired_entry.get("sha256")
and file_path
and os.path.exists(file_path)
):
try:
sha256 = await calculate_sha256(file_path)
if sha256:
desired_entry["sha256"] = sha256.lower()
except Exception:
pass
# Not in cache at all — delegate to the full update path ------------
if existing_entry is None:
result = await self.update_single_model_cache(
file_path, file_path, metadata_dict
)
return bool(result)
# Compare — skip everything if already in sync -----------------------
if not self._cache_entries_differ(existing_entry, desired_entry):
return False
# Re-validate: the cache may have been replaced concurrently
# (e.g. by _apply_scan_result). Use identity check, not equality,
# so we detect when the raw_data list was swapped out from under us.
if self._cache is None or not any(
item is existing_entry for item in self._cache.raw_data
):
return False
# ---- Differences detected: apply targeted, in-place updates --------
# Snapshot old values for delta computations
old_tags = list(existing_entry.get("tags") or [])
old_sha256: str = existing_entry.get("sha256", "") or ""
old_model_name: str = existing_entry.get("model_name", "") or ""
old_file_name: str = existing_entry.get("file_name", "") or ""
old_modified: float = float(existing_entry.get("modified", 0.0) or 0.0)
old_size: int = int(existing_entry.get("size", 0) or 0)
old_civitai = existing_entry.get("civitai")
# ---- In-place update of the cache entry ----
existing_entry.clear()
existing_entry.update(desired_entry)
# ---- Incremental tag count update ----
new_tags: set = set(desired_entry.get("tags") or [])
old_tag_set: set = set(old_tags)
for tag in old_tag_set - new_tags:
current = self._tags_count.get(tag, 0)
if current <= 1:
self._tags_count.pop(tag, None)
else:
self._tags_count[tag] = current - 1
for tag in new_tags - old_tag_set:
self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
# ---- Incremental hash index update ----
new_sha = (desired_entry.get("sha256", "") or "").lower()
old_sha = (old_sha256 or "").lower()
if new_sha != old_sha:
if old_sha:
self._hash_index.remove_by_path(file_path)
if new_sha:
self._hash_index.add_entry(new_sha, file_path)
# ---- Incremental version index update ----
new_civitai = desired_entry.get("civitai")
if old_civitai != new_civitai:
temp_old = {
"file_path": file_path,
"file_name": old_file_name,
"civitai": old_civitai,
}
cache.remove_from_version_index(temp_old)
cache.add_to_version_index(existing_entry)
# ---- Conditional resort (only when sort-key fields changed) ----
need_resort = False
_last = cache._last_sort
sort_key: Optional[str] = _last[0] if _last != (None, None) else None
if sort_key == "name":
if (
old_model_name != desired_entry.get("model_name", "")
or old_file_name != desired_entry.get("file_name", "")
):
need_resort = True
elif sort_key == "date":
if old_modified != float(desired_entry.get("modified", 0.0) or 0.0):
need_resort = True
elif sort_key == "size":
if old_size != int(desired_entry.get("size", 0) or 0):
need_resort = True
if need_resort:
await cache.resort()
# ---- Targeted SQL update (single row, not full save_cache) ----
persistent = getattr(self, "_persistent_cache", None)
if persistent is not None:
old_item_for_sql: Dict[str, Any] = {
"file_path": file_path,
"tags": old_tags,
"sha256": old_sha256,
}
await asyncio.get_event_loop().run_in_executor(
None,
persistent.update_single_model,
self.model_type,
desired_entry,
old_item_for_sql,
)
return True
@staticmethod
def _cache_entries_differ(a: Dict[str, Any], b: Dict[str, Any]) -> bool:
"""Return ``True`` when two cache-entry dicts differ in any field.
Tag lists are compared order-insensitively; all other keys use
standard equality.
"""
a_tags = sorted(a.get("tags") or [])
b_tags = sorted(b.get("tags") or [])
if a_tags != b_tags:
return True
all_keys = set(a.keys()) | set(b.keys())
for key in all_keys:
if key == "tags":
continue
if a.get(key) != b.get(key):
return True
return False
def has_hash(self, sha256: str) -> bool:
"""Check if a model with given hash exists"""
return self._hash_index.has_hash(sha256.lower())
@@ -1580,7 +1859,32 @@ class ModelScanner:
if limit == 0:
return sorted_tags
return sorted_tags[:limit]
async def search_tags(
self, query: str, limit: int = 50
) -> List[Dict[str, any]]:
"""Search tags by case-insensitive substring match, sorted by count.
If query is empty, behaves like get_top_tags (returns top ``limit``
tags). If limit is 0, all matching tags are returned.
"""
await self.get_cached_data()
normalized_query = (query or "").strip().lower()
if not normalized_query:
return await self.get_top_tags(limit if limit > 0 else 20)
matched = [
{"tag": tag, "count": count}
for tag, count in self._tags_count.items()
if normalized_query in tag.lower()
]
matched.sort(key=lambda x: x["count"], reverse=True)
if limit == 0:
return matched
return matched[:limit]
async def get_base_models(self, limit: int = 20) -> List[Dict[str, any]]:
"""Get base models sorted by count. If limit is 0, return all."""
cache = await self.get_cached_data()
@@ -1696,6 +2000,8 @@ class ModelScanner:
break
try:
_require_path_in_library_roots(file_path, self, label="File path")
target_dir = os.path.dirname(file_path)
base_name = os.path.basename(file_path)
file_name, main_extension = os.path.splitext(base_name)
@@ -1761,6 +2067,13 @@ class ModelScanner:
"""
if not file_paths or self._cache is None:
return False
if self.is_cancelled():
logger.info(
f"{self.model_type.capitalize()} Scanner: Skipping cache update "
"after cancelled bulk delete"
)
return False
try:
# Get all models that need to be removed from cache
+49 -2
View File
@@ -724,6 +724,16 @@ class ModelUpdateService:
"Refreshing update metadata for %d %s models", total_models, model_type
)
# When filtering by folder, also collect the cross-folder version set
# so that versions already present in other folders are not reported
# as available updates. See issue #997.
all_local_versions: Optional[Dict[int, List[int]]] = None
if folder_path is not None:
all_local_versions = await self._collect_local_versions(
scanner,
target_model_ids=target_filter,
)
results: Dict[int, ModelUpdateRecord] = {}
prefetched: Dict[int, Mapping] = {}
@@ -762,6 +772,12 @@ class ModelUpdateService:
for index, (model_id, version_ids) in enumerate(
local_versions.items(), start=1
):
# Use cross-folder version IDs for is_in_library if available
all_vids: Sequence[int] = (
all_local_versions.get(model_id, [])
if all_local_versions is not None
else version_ids
)
record = await self._refresh_single_model(
model_type,
model_id,
@@ -769,6 +785,7 @@ class ModelUpdateService:
metadata_provider,
force_refresh=force_refresh,
prefetched_response=prefetched.get(model_id),
all_local_version_ids=all_vids,
)
if scanner.is_cancelled():
logger.info(f"{model_type.capitalize()} Update Service: Refresh cancelled by user")
@@ -964,8 +981,16 @@ class ModelUpdateService:
*,
force_refresh: bool = False,
prefetched_response: Optional[Mapping] = None,
all_local_version_ids: Optional[Sequence[int]] = None,
) -> Optional[ModelUpdateRecord]:
normalized_local = self._normalize_sequence(local_versions)
# When folder-filtering, this carries the cross-folder version set
# for is_in_library; otherwise it falls back to normalized_local.
normalized_all = (
self._normalize_sequence(all_local_version_ids)
if all_local_version_ids is not None
else normalized_local
)
now = time.time()
async with self._lock:
existing = self._get_record(model_type, model_id)
@@ -973,6 +998,7 @@ class ModelUpdateService:
record = self._merge_with_local_versions(
existing,
normalized_local,
all_local_version_ids=normalized_all,
)
self._upsert_record(record)
return record
@@ -1048,6 +1074,7 @@ class ModelUpdateService:
record = self._merge_with_local_versions(
existing,
normalized_local,
all_local_version_ids=normalized_all,
)
self._upsert_record(record)
return record
@@ -1059,6 +1086,7 @@ class ModelUpdateService:
model_type=model_type,
model_id=model_id,
last_checked_at=now,
all_local_version_ids=normalized_all,
)
record = replace(record, should_ignore_model=True)
self._upsert_record(record)
@@ -1077,6 +1105,7 @@ class ModelUpdateService:
fetched_versions,
existing,
now,
all_local_version_ids=normalized_all,
)
else:
record = self._merge_with_local_versions(
@@ -1085,6 +1114,7 @@ class ModelUpdateService:
model_type=model_type,
model_id=model_id,
last_checked_at=existing.last_checked_at if existing else None,
all_local_version_ids=normalized_all,
)
self._upsert_record(record)
return record
@@ -1322,12 +1352,20 @@ class ModelUpdateService:
existing: Optional[ModelUpdateRecord],
normalized_local: Sequence[int],
*,
all_local_version_ids: Optional[Sequence[int]] = None,
model_type: Optional[str] = None,
model_id: Optional[int] = None,
last_checked_at: Optional[float] = None,
version_info: Optional[Mapping] = None,
) -> ModelUpdateRecord:
local_set = set(normalized_local)
# When folder-filtering, also consider versions in other folders
# as in-library so they are not reported as available updates.
effective_local_set: set[int] = (
local_set | set(all_local_version_ids)
if all_local_version_ids is not None
else local_set
)
versions: List[ModelVersionRecord] = []
ignore_map: Dict[int, bool] = {}
if existing:
@@ -1339,7 +1377,7 @@ class ModelUpdateService:
versions.append(
replace(
version,
is_in_library=version.version_id in local_set,
is_in_library=version.version_id in effective_local_set,
)
)
elif model_type is None or model_id is None:
@@ -1386,8 +1424,17 @@ class ModelUpdateService:
remote_versions: Sequence[ModelVersionRecord],
existing: Optional[ModelUpdateRecord],
timestamp: float,
*,
all_local_version_ids: Optional[Sequence[int]] = None,
) -> ModelUpdateRecord:
local_set = set(local_versions)
# When folder-filtering, also consider versions in other folders
# as in-library so they are not reported as available updates.
effective_local_set: set[int] = (
local_set | set(all_local_version_ids)
if all_local_version_ids is not None
else local_set
)
ignore_map = {version.version_id: version.should_ignore for version in existing.versions} if existing else {}
preview_map = {version.version_id: version.preview_url for version in existing.versions} if existing else {}
sort_map = {version.version_id: version.sort_index for version in existing.versions} if existing else {}
@@ -1406,7 +1453,7 @@ class ModelUpdateService:
released_at=remote_version.released_at,
size_bytes=remote_version.size_bytes,
preview_url=remote_version.preview_url or preview_map.get(version_id),
is_in_library=version_id in local_set,
is_in_library=version_id in effective_local_set,
should_ignore=ignore_map.get(version_id, remote_version.should_ignore),
sort_index=sort_map.get(version_id, index),
early_access_ends_at=remote_version.early_access_ends_at,
+103 -9
View File
@@ -57,6 +57,7 @@ class PersistentModelCache:
"db_checked",
"last_checked_at",
"hash_status",
"hf_url",
)
_MODEL_UPDATE_COLUMNS: Tuple[str, ...] = _MODEL_COLUMNS[2:]
_instances: Dict[str, "PersistentModelCache"] = {}
@@ -165,8 +166,8 @@ class PersistentModelCache:
item = {
"file_path": file_path,
"file_name": row["file_name"],
"model_name": row["model_name"],
"file_name": row["file_name"] or "",
"model_name": row["model_name"] or "",
"folder": row["folder"] or "",
"size": row["size"] or 0,
"modified": row["modified"] or 0.0,
@@ -188,6 +189,7 @@ class PersistentModelCache:
"skip_metadata_refresh": bool(row["skip_metadata_refresh"]),
"license_flags": int(license_value),
"hash_status": row["hash_status"] or "completed",
"hf_url": row["hf_url"] or "",
}
raw_data.append(item)
@@ -452,6 +454,7 @@ class PersistentModelCache:
db_checked INTEGER,
last_checked_at REAL,
hash_status TEXT,
hf_url TEXT DEFAULT '',
PRIMARY KEY (model_type, file_path)
);
@@ -500,6 +503,7 @@ class PersistentModelCache:
# Persisting without explicit flags should assume CivitAI's documented defaults (0b111001 == 57).
"license_flags": f"INTEGER DEFAULT {DEFAULT_LICENSE_FLAGS}",
"hash_status": "TEXT DEFAULT 'completed'",
"hf_url": "TEXT DEFAULT ''",
}
for column, definition in required_columns.items():
@@ -548,19 +552,19 @@ class PersistentModelCache:
return (
model_type,
item.get("file_path"),
item.get("file_name"),
item.get("model_name"),
item.get("folder"),
item.get("file_name") or "",
item.get("model_name") or "",
item.get("folder") or "",
int(item.get("size") or 0),
float(item.get("modified") or 0.0),
(item.get("sha256") or "").lower() or None,
item.get("base_model"),
item.get("preview_url"),
item.get("base_model") or "",
item.get("preview_url") or "",
int(item.get("preview_nsfw_level") or 0),
1 if item.get("from_civitai", True) else 0,
1 if item.get("favorite") else 0,
item.get("notes"),
item.get("usage_tips"),
item.get("notes") or "",
item.get("usage_tips") or "",
metadata_source,
civitai.get("id"),
civitai.get("modelId"),
@@ -575,6 +579,7 @@ class PersistentModelCache:
1 if item.get("db_checked") else 0,
float(item.get("last_checked_at") or 0.0),
item.get("hash_status", "completed"),
item.get("hf_url") or "",
)
def _insert_model_sql(self) -> str:
@@ -582,6 +587,95 @@ class PersistentModelCache:
placeholders = ", ".join(["?"] * len(self._MODEL_COLUMNS))
return f"INSERT INTO models ({columns}) VALUES ({placeholders})"
def update_single_model(
self,
model_type: str,
new_item: Dict,
old_item: Optional[Dict] = None,
) -> None:
"""Update a single model row in the persistent cache.
A lightweight alternative to :meth:`save_cache` that performs a targeted
DELETE + INSERT for the model row and computes incremental tag / hash-index
deltas from *old_item*. When *old_item* is omitted the previous tags and
hash are not cleaned up (callers should only omit it for brand-new entries).
All operations run inside a single transaction so readers see a consistent
view.
"""
if not self.is_enabled():
return
if not self._schema_initialized:
self._initialize_schema()
if not self._schema_initialized:
return
file_path: Optional[str] = new_item.get("file_path")
if not file_path:
return
try:
with self._db_lock:
conn = self._connect()
try:
conn.execute("PRAGMA foreign_keys = ON")
conn.execute("BEGIN")
# --- model row (DELETE + INSERT = upsert) ---
conn.execute(
"DELETE FROM models WHERE model_type = ? AND file_path = ?",
(model_type, file_path),
)
row = self._prepare_model_row(model_type, new_item)
conn.execute(self._insert_model_sql(), row)
# --- tags ---
new_tags: set = set(new_item.get("tags") or [])
old_tags: set = set(old_item.get("tags") or []) if old_item else set()
tags_to_delete = old_tags - new_tags
tags_to_insert = new_tags - old_tags
if tags_to_delete:
conn.executemany(
"DELETE FROM model_tags WHERE model_type = ? AND file_path = ? AND tag = ?",
[(model_type, file_path, t) for t in tags_to_delete],
)
if tags_to_insert:
conn.executemany(
"INSERT INTO model_tags (model_type, file_path, tag) VALUES (?, ?, ?)",
[(model_type, file_path, t) for t in tags_to_insert],
)
# --- hash_index ---
new_sha: Optional[str] = (new_item.get("sha256") or "").lower() or None
old_sha: Optional[str] = (
(old_item.get("sha256") or "").lower() or None
) if old_item else None
if new_sha != old_sha:
if old_sha:
conn.execute(
"DELETE FROM hash_index WHERE model_type = ? AND sha256 = ? AND file_path = ?",
(model_type, old_sha, file_path),
)
if new_sha:
conn.execute(
"INSERT OR IGNORE INTO hash_index (model_type, sha256, file_path) VALUES (?, ?, ?)",
(model_type, new_sha, file_path),
)
conn.execute("COMMIT")
except Exception:
conn.execute("ROLLBACK")
raise
finally:
conn.close()
except Exception as exc:
logger.warning(
"Failed to update single model in persistent cache (%s): %s",
file_path,
exc,
)
def _load_tags(self, conn: sqlite3.Connection, model_type: str) -> Dict[str, List[str]]:
tag_rows = conn.execute(
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
+6 -2
View File
@@ -1,7 +1,6 @@
import asyncio
from typing import Iterable, List, Dict, Optional
from dataclasses import dataclass, field
from operator import itemgetter
from natsort import natsorted
@@ -149,5 +148,10 @@ class RecipeCache:
)
if not name_only:
self.sorted_by_date = sorted(
self.raw_data, key=itemgetter("created_date", "file_path"), reverse=True
self.raw_data,
key=lambda x: (
x.get("modified", x.get("created_date", 0)),
x.get("file_path", ""),
),
reverse=True,
)
+16 -48
View File
@@ -21,7 +21,7 @@ from .checkpoint_scanner import CheckpointScanner
from .settings_manager import get_settings_manager
from .recipes.errors import RecipeNotFoundError
from ..utils.civitai_utils import extract_civitai_image_id
from ..utils.utils import calculate_recipe_fingerprint, fuzzy_match
from ..utils.utils import calculate_recipe_fingerprint
from natsort import natsorted
import sys
import re
@@ -1020,13 +1020,16 @@ class RecipeScanner:
try:
result = self._fts_index.search(search, fields)
# Return None if empty to trigger fuzzy fallback
# Empty FTS results may indicate query syntax issues or need for fuzzy matching
# Return empty set for empty FTS results — do NOT fall back to
# Python fuzzy matching, which freezes the server with 10k+ recipes.
# FTS5 prefix matching with unicode61 tokenizer correctly handles
# compound tokens (e.g. "illustrious" matches "path/illustrious/model").
# If FTS returns nothing, there are genuinely no matching recipes.
if not result:
return None
return set()
return result
except Exception as exc:
logger.debug("FTS search failed, falling back to fuzzy search: %s", exc)
logger.debug("FTS search failed, falling back to title-only search: %s", exc)
return None
def _update_fts_index_for_recipe(
@@ -2079,49 +2082,14 @@ class RecipeScanner:
if str(item.get("id", "")) in fts_matching_ids
]
else:
# Fallback to fuzzy_match (slower but always available)
# Build the search predicate based on search options
def matches_search(item):
# Search in title if enabled
if search_options.get("title", True):
if fuzzy_match(str(item.get("title", "")), search):
return True
# Search in tags if enabled
if search_options.get("tags", True) and "tags" in item:
for tag in item["tags"]:
if fuzzy_match(tag, search):
return True
# Search in lora file names if enabled
if search_options.get("lora_name", True) and "loras" in item:
for lora in item["loras"]:
if fuzzy_match(str(lora.get("file_name", "")), search):
return True
# Search in lora model names if enabled
if search_options.get("lora_model", True) and "loras" in item:
for lora in item["loras"]:
if fuzzy_match(str(lora.get("modelName", "")), search):
return True
# Search in prompt and negative_prompt if enabled
if search_options.get("prompt", True) and "gen_params" in item:
gen_params = item["gen_params"]
if fuzzy_match(str(gen_params.get("prompt", "")), search):
return True
if fuzzy_match(
str(gen_params.get("negative_prompt", "")), search
):
return True
# No match found
return False
# Filter the data using the search predicate
filtered_data = [
item for item in filtered_data if matches_search(item)
]
# FTS index not yet built — return empty rather than
# scanning 42k+ items in Python. The FTS background build
# finishes in seconds; by the time a user navigates here
# and types a search, it is already available.
logger.debug(
"FTS index not ready — search '%s' returning empty", search
)
filtered_data = []
# Apply additional filters
if filters:
+96 -17
View File
@@ -146,11 +146,38 @@ class RecipeAnalysisService:
):
metadata = metadata["meta"]
# Include modelVersionIds from root level if available
# Civitai API returns modelVersionIds at root level, not in meta
# Include modelVersionIds from root level if available.
# CivitAI API returns modelVersionIds at root level, not in meta.
# When meta is null (None), create a minimal dict so downstream
# parsers can still discover LoRAs and checkpoints.
model_version_ids = image_info.get("modelVersionIds")
if model_version_ids and isinstance(metadata, dict):
metadata["modelVersionIds"] = model_version_ids
if model_version_ids:
if isinstance(metadata, dict):
metadata["modelVersionIds"] = model_version_ids
else:
metadata = {"modelVersionIds": model_version_ids}
# Inject browsingLevel (canonical integer) so the recipe's
# preview_nsfw_level can be set, enabling proper NSFW blur
# of the preview image. Fall back to nsfwLevel (string)
# when browsingLevel is absent.
if isinstance(metadata, dict):
browsing_level = image_info.get("browsingLevel")
nsfw_level_str = image_info.get("nsfwLevel")
if isinstance(browsing_level, int) and browsing_level > 0:
metadata["browsingLevel"] = browsing_level
elif (
isinstance(nsfw_level_str, str)
and nsfw_level_str
in (
"PG", "PG13", "R", "X", "XXX", "Blocked",
)
):
from ...utils.constants import NSFW_LEVELS
metadata["browsingLevel"] = NSFW_LEVELS.get(
nsfw_level_str, 0
)
# Validate that metadata contains meaningful recipe fields
# If not, treat as None to trigger EXIF extraction from downloaded image
@@ -171,12 +198,19 @@ class RecipeAnalysisService:
temp_path = self._create_temp_path(suffix=extension)
await self._download_image(url, temp_path)
if metadata is None and not is_video:
metadata = await asyncio.to_thread(
# Always extract EXIF from the downloaded image for generation
# params (prompt, negative prompt, sampler, steps, etc.).
# Previously this was gated on ``metadata is None``, but that
# skipped EXIF entirely when API metadata (modelVersionIds,
# browsingLevel) is present, losing all generation parameters.
exif_metadata = None
if not is_video:
exif_metadata = await asyncio.to_thread(
self._exif_utils.extract_image_metadata, temp_path
)
if not metadata and civitai_image_id and image_info:
# Fallback: try the original (non-optimized) image for EXIF data
if not exif_metadata and civitai_image_id and image_info:
original_url = image_info.get("url")
if original_url:
self._logger.debug(
@@ -187,15 +221,38 @@ class RecipeAnalysisService:
orig_temp_path = self._create_temp_path(suffix=".png")
try:
await self._download_image(original_url, orig_temp_path)
metadata = await asyncio.to_thread(
exif_metadata = await asyncio.to_thread(
self._exif_utils.extract_image_metadata,
orig_temp_path,
)
finally:
self._safe_cleanup(orig_temp_path)
# Parse EXIF data (typically a string like parameters/prompt/workflow)
# and API metadata (dict with modelVersionIds, browsingLevel) separately,
# then merge: API loras/checkpoint override, EXIF gen_params fill in gaps.
# This mirrors the two-pass approach in _do_import_from_url.
exif_parsed_result = None
if isinstance(exif_metadata, str):
exif_parser = self._recipe_parser_factory.create_parser(exif_metadata)
if exif_parser:
exif_data = await exif_parser.parse_metadata(
exif_metadata, recipe_scanner=recipe_scanner,
)
if exif_data and not exif_data.get("error"):
exif_parsed_result = exif_data
# Merge API metadata (dict) with EXIF data (if dict) for the
# CivitaiApiMetadataParser. If EXIF data is a string it was
# parsed above — don't try to merge a string into a dict.
merged = {}
if isinstance(exif_metadata, dict):
merged.update(exif_metadata)
if isinstance(metadata, dict):
merged.update(metadata)
result = await self._parse_metadata(
metadata or {},
merged,
recipe_scanner=recipe_scanner,
image_path=temp_path,
include_image_base64=True,
@@ -203,13 +260,23 @@ class RecipeAnalysisService:
extension=extension,
)
if civitai_image_id and image_info and not result.payload.get("error"):
mvid = image_info.get("modelVersionId")
if not mvid:
mvids = image_info.get("modelVersionIds")
if isinstance(mvids, list) and mvids:
mvid = mvids[0]
# Merge EXIF string-parsed gen_params into the API result.
# API gen_params take priority (they come later via update).
if exif_parsed_result and not result.payload.get("error"):
exif_gp = exif_parsed_result.get("gen_params") or {}
result_gp = result.payload.get("gen_params") or {}
merged_gp = {**exif_gp, **result_gp}
if merged_gp:
result.payload["gen_params"] = merged_gp
if civitai_image_id and image_info and not result.payload.get("error"):
# Use the metadata dict we built (may contain modelVersionIds
# and browsingLevel from the API root level). Do NOT pass
# image_info.get("meta") — it is null for images whose meta
# lives at the root level only. Also do NOT derive
# model_version_id from modelVersionIds[0] — that array mixes
# checkpoints, LoRAs, and other types without ordering
# guarantees; the parser already resolved them correctly.
recipe_for_enrich = {
"gen_params": result.payload.get("gen_params", {}),
"loras": result.payload.get("loras", []),
@@ -222,8 +289,10 @@ class RecipeAnalysisService:
recipe=recipe_for_enrich,
civitai_client=civitai_client,
request_params=None,
prefetched_civitai_meta_raw=image_info.get("meta"),
prefetched_model_version_id=mvid,
prefetched_civitai_meta_raw=(
metadata if isinstance(metadata, dict) else None
),
prefetched_model_version_id=None,
)
result.payload["gen_params"] = recipe_for_enrich["gen_params"]
@@ -232,6 +301,12 @@ class RecipeAnalysisService:
if recipe_for_enrich.get("base_model"):
result.payload["base_model"] = recipe_for_enrich["base_model"]
# Extract browsingLevel from our constructed metadata for NSFW blur
if isinstance(metadata, dict):
bl = metadata.get("browsingLevel")
if isinstance(bl, int) and bl > 0:
result.payload["preview_nsfw_level"] = bl
return result
finally:
if temp_path:
@@ -314,6 +389,10 @@ class RecipeAnalysisService:
"prompt_type",
"positive",
"negative",
# modelVersionIds is injected at the root level by CivitAI's image
# API when meta is null. It carries the version IDs of ALL models
# (checkpoint + LoRAs) used to generate the image.
"modelVersionIds",
}
return any(field in metadata for field in recipe_fields)
+2 -1
View File
@@ -216,11 +216,12 @@ class RecipePersistenceService:
"preview_nsfw_level",
"favorite",
"gen_params",
"base_model",
)
if not any(key in updates for key in allowed_fields):
raise RecipeValidationError(
"At least one field to update must be provided (title or tags or source_path or preview_nsfw_level or favorite or gen_params)"
"At least one field to update must be provided (title or tags or source_path or preview_nsfw_level or favorite or gen_params or base_model)"
)
if "gen_params" in updates and not isinstance(updates["gen_params"], dict):
+133 -39
View File
@@ -65,6 +65,8 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
"onboarding_completed": False,
"dismissed_banners": [],
"enable_metadata_archive_db": False,
"enable_civarchive_api": True,
"metadata_provider_order": "civitai_archive_sqlite",
"proxy_enabled": False,
"proxy_host": "",
"proxy_port": "",
@@ -91,7 +93,6 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
"autoplay_on_hover": False,
"display_density": "default",
"card_info_display": "always",
"show_folder_sidebar": True,
"include_trigger_words": False,
"compact_mode": False,
"priority_tags": DEFAULT_PRIORITY_TAG_CONFIG.copy(),
@@ -99,13 +100,20 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
"lora_syntax_format": "legacy",
"model_card_footer_action": "replace_preview",
"show_version_on_card": True,
"update_flag_strategy": "same_base",
"version_grouping": "same_base",
"auto_organize_exclusions": [],
"metadata_refresh_skip_paths": [],
"skip_previously_downloaded_model_versions": False,
"download_skip_base_models": [],
"backup_auto_enabled": True,
"backup_retention_count": 5,
"use_new_license_icons": True,
"group_by_model": False,
# AI / LLM provider configuration (BYOK)
"llm_provider": "openai", # "openai" | "ollama" | "custom"
"llm_api_key": "",
"llm_api_base": "", # empty = provider default
"llm_model": "", # e.g. "gpt-4o-mini"
}
@@ -134,6 +142,9 @@ class SettingsManager:
self._template_path = (
Path(__file__).resolve().parents[2] / "settings.json.example"
)
# Known placeholder value in settings.json.example; any file containing
# this value should be treated as "not configured".
self._TEMPLATE_PLACEHOLDER_API_KEY = "your_civitai_api_key_here"
self.settings = self._load_settings()
self._migrate_setting_keys()
self._ensure_default_settings()
@@ -143,6 +154,11 @@ class SettingsManager:
self._check_environment_variables()
self._collect_configuration_warnings()
if os.environ.get("LORA_MANAGER_PORTABLE", "0") == "1":
if not self.settings.get("use_portable_settings"):
self.settings["use_portable_settings"] = True
self._save_settings()
if self._needs_initial_save:
self._save_settings()
self._needs_initial_save = False
@@ -165,6 +181,12 @@ class SettingsManager:
self._original_disk_payload = copy.deepcopy(data)
if self._matches_template_payload(data):
self._preserve_disk_template = True
# Clean up the template placeholder so it is not treated
# as a real key (affects both the frontend boolean and
# the downloader's Authorization header).
placeholder = self._TEMPLATE_PLACEHOLDER_API_KEY
if data.get("civitai_api_key") == placeholder:
data["civitai_api_key"] = ""
return data
except json.JSONDecodeError as exc:
logger.error("Failed to parse settings.json: %s", exc)
@@ -610,12 +632,37 @@ class SettingsManager:
return False
@staticmethod
def _normalize_path_set(paths: Iterable[str]) -> set[str]:
"""Normalize an iterable of paths for set-based overlap comparison.
Resolves symlinks via ``os.path.realpath`` when the path exists on disk,
then applies ``os.path.normcase`` + ``os.path.normpath`` for consistent
cross-platform comparison. Non-string / empty entries are skipped.
"""
result: set[str] = set()
for p in paths:
if not isinstance(p, str):
continue
stripped = p.strip()
if not stripped:
continue
if os.path.exists(stripped):
stripped = os.path.normpath(os.path.realpath(stripped))
result.add(os.path.normcase(stripped))
return result
def _validate_folder_paths(
self,
library_name: str,
folder_paths: Mapping[str, Iterable[str]],
) -> None:
"""Ensure folder paths do not overlap with other libraries."""
"""Ensure folder paths do not overlap with other libraries.
Also detects checkpoints unet path overlap within the same library
(including via symlink resolution), which is a configuration error since
these model types must use separate physical folders.
"""
libraries = self.settings.get("libraries", {})
normalized_new: Dict[str, Dict[str, str]] = {}
for key, values in folder_paths.items():
@@ -653,6 +700,22 @@ class SettingsManager:
f"Folder path(s) {collisions} already assigned to library '{other_name}'"
)
# Checkpoints ↔ unet overlap within the same library
ckpt_paths = folder_paths.get("checkpoints", []) or []
unet_paths = folder_paths.get("unet", []) or []
if ckpt_paths and unet_paths:
ckpt_real = self._normalize_path_set(ckpt_paths)
unet_real = self._normalize_path_set(unet_paths)
overlap = ckpt_real & unet_real
if overlap:
collisions = ", ".join(sorted(overlap))
raise ValueError(
f"Path(s) {collisions} are configured for both "
f"'checkpoints' and 'unet' (diffusion models). "
f"These model types must use separate physical folders. "
f"Please remove one of the conflicting entries."
)
def _update_active_library_entry(
self,
*,
@@ -735,6 +798,7 @@ class SettingsManager:
"includeTriggerWords": "include_trigger_words",
"compactMode": "compact_mode",
"modelCardFooterAction": "model_card_footer_action",
"update_flag_strategy": "version_grouping",
}
updated = False
@@ -862,6 +926,23 @@ class SettingsManager:
self.settings["civitai_api_key"] = env_api_key
self._save_settings()
# LLM provider overrides
llm_env_map = {
"LLM_API_KEY": "llm_api_key",
"LLM_MODEL": "llm_model",
"LLM_API_BASE": "llm_api_base",
"LLM_PROVIDER": "llm_provider",
}
llm_changed = False
for env_var, settings_key in llm_env_map.items():
env_val = os.environ.get(env_var)
if env_val:
logger.info("Found %s environment variable", env_var)
self.settings[settings_key] = env_val
llm_changed = True
if llm_changed:
self._save_settings()
def _default_settings_actions(self) -> List[Dict[str, Any]]:
return [
{
@@ -1509,8 +1590,12 @@ class SettingsManager:
portable_switch_pending = True
self._prepare_portable_switch(value)
if key == "folder_paths" and isinstance(value, Mapping):
active_name = self.get_active_library_name()
self._validate_folder_paths(active_name, value)
self._update_active_library_entry(folder_paths=value) # type: ignore[arg-type]
elif key == "extra_folder_paths" and isinstance(value, Mapping):
active_name = self.get_active_library_name()
self._validate_folder_paths(active_name, value)
self._update_active_library_entry(extra_folder_paths=value) # type: ignore[arg-type]
elif key == "default_lora_root":
self._update_active_library_entry(default_lora_root=str(value))
@@ -1557,7 +1642,7 @@ class SettingsManager:
previous_dir = os.path.dirname(previous_path) or target_dir
if os.path.abspath(previous_path) != os.path.abspath(target_path):
self._copy_model_cache_directory(previous_dir, target_dir)
self._migrate_settings_directory_content(previous_dir, target_dir)
logger.info("Switching settings file to: %s", target_path)
self._pending_portable_switch = {"other_path": other_path}
@@ -1592,46 +1677,52 @@ class SettingsManager:
finally:
self._pending_portable_switch = None
def _copy_model_cache_directory(self, source_dir: str, target_dir: str) -> None:
"""Copy model_cache artifacts when switching storage locations."""
def _migrate_settings_directory_content(
self, source_dir: str, target_dir: str
) -> None:
"""Migrate settings directory subdirectories when switching storage locations.
Copies the canonical subdirectories (cache, backups, logs, stats, wildcards)
from the old settings directory to the new one. Legacy cache artifacts
(model_cache, recipe_cache, etc.) are migrated lazily by
``resolve_cache_path_with_migration`` on first access.
Args:
source_dir: The previous settings directory path.
target_dir: The new settings directory path.
"""
if not source_dir or not target_dir:
return
source_cache_dir = os.path.join(source_dir, "model_cache")
target_cache_dir = os.path.join(target_dir, "model_cache")
if os.path.isdir(source_cache_dir) and os.path.abspath(
source_cache_dir
) != os.path.abspath(target_cache_dir):
try:
shutil.copytree(
source_cache_dir,
target_cache_dir,
dirs_exist_ok=True,
ignore=shutil.ignore_patterns("*.sqlite-shm", "*.sqlite-wal"),
)
except Exception as exc:
logger.warning(
"Failed to copy model_cache directory from %s to %s: %s",
source_cache_dir,
target_cache_dir,
exc,
)
def _copy_dir(name: str) -> None:
source = os.path.join(source_dir, name)
target = os.path.join(target_dir, name)
if os.path.isdir(source) and os.path.abspath(source) != os.path.abspath(
target
):
try:
shutil.copytree(
source,
target,
dirs_exist_ok=True,
ignore=shutil.ignore_patterns("*.sqlite-shm", "*.sqlite-wal"),
)
except Exception as exc:
logger.warning(
"Failed to copy directory %s from %s to %s: %s",
name,
source,
target,
exc,
)
source_cache_file = os.path.join(source_dir, "model_cache.sqlite")
target_cache_file = os.path.join(target_dir, "model_cache.sqlite")
if os.path.isfile(source_cache_file) and os.path.abspath(
source_cache_file
) != os.path.abspath(target_cache_file):
try:
shutil.copy2(source_cache_file, target_cache_file)
except Exception as exc:
logger.warning(
"Failed to copy model_cache.sqlite from %s to %s: %s",
source_cache_file,
target_cache_file,
exc,
)
# Managed subdirectories under settings_dir
_copy_dir("cache")
_copy_dir("backups")
_copy_dir("logs")
_copy_dir("stats")
_copy_dir("wildcards")
def _get_user_config_directory(self) -> str:
"""Return the user configuration directory, falling back to ~/.config."""
@@ -1758,6 +1849,9 @@ class SettingsManager:
if key in self.settings:
minimal[key] = copy.deepcopy(self.settings[key])
if self.settings.get("use_portable_settings"):
minimal["use_portable_settings"] = True
if self._seed_template:
for key, value in self._seed_template.items():
minimal.setdefault(key, copy.deepcopy(value))
+2 -2
View File
@@ -36,9 +36,9 @@ class TagUpdateService:
if isinstance(tag, str) and tag.strip():
# Convert all tags to lowercase to avoid case sensitivity issues on Windows
normalized = tag.strip().lower()
if normalized.lower() not in existing_lower:
if normalized not in existing_lower:
existing_tags.append(normalized)
existing_lower.append(normalized.lower())
existing_lower.append(normalized)
tags_added.append(normalized)
metadata["tags"] = existing_tags
@@ -3,6 +3,7 @@
from __future__ import annotations
import logging
import time
from typing import Any, Dict, List, Optional, Protocol, Sequence
from ..metadata_sync_service import MetadataSyncService
@@ -50,6 +51,10 @@ class BulkMetadataRefreshUseCase:
if not model.get("skip_metadata_refresh", False)
and not self._is_in_skip_path(model.get("folder", ""), skip_paths)
and (not model.get("civitai") or not model["civitai"].get("id"))
# Skip models downloaded from Hugging Face — they are not on
# CivitAI / CivArchive. Users can still refresh them individually
# via the right-click context menu.
and not model.get("hf_url", "")
and not (
# Skip models confirmed not on CivitAI when no need to retry
model.get("from_civitai") is False
@@ -62,26 +67,48 @@ class BulkMetadataRefreshUseCase:
]
total_to_process = len(to_process)
initial_skipped = total_models - total_to_process # models excluded from fetch queue
processed = 0
success = 0
skipped_count = initial_skipped
handled_count = initial_skipped
needs_resort = False
start_time = time.monotonic()
failures: List[Dict[str, str]] = []
self._service.scanner.reset_cancellation()
async def emit(status: str, **extra: Any) -> None:
if progress_callback is None:
return
payload = {"status": status, "total": total_to_process, "processed": processed, "success": success}
payload = {
"status": status,
"total": total_models,
"processed": processed,
"success": success,
"failure_count": len(failures),
"skipped_count": skipped_count,
"handled": handled_count,
"elapsed_seconds": int(time.monotonic() - start_time),
}
# Only include full failure details in terminal emits (completed,
# cancelled, rate_limited) to avoid serializing the list on every
# per-model progress update.
if failures and status in ("completed", "cancelled", "rate_limited"):
payload["failures"] = failures
payload.update(extra)
await progress_callback.on_progress(payload)
await emit("started")
RATE_LIMIT_ABORT_THRESHOLD = 3
consecutive_rate_limits = 0
for model in to_process:
if self._service.scanner.is_cancelled():
self._logger.info("Bulk metadata refresh cancelled by user")
await emit("cancelled", processed=processed, success=success)
return {"success": False, "message": "Operation cancelled", "processed": processed, "updated": success, "total": total_models}
return {"success": False, "message": "Operation cancelled", "processed": processed, "updated": success, "total": total_models, "failures": failures, "failure_count": len(failures), "skipped_count": skipped_count, "elapsed_seconds": int(time.monotonic() - start_time)}
try:
original_name = model.get("model_name")
@@ -99,33 +126,89 @@ class BulkMetadataRefreshUseCase:
if sha256:
model["sha256"] = sha256
model["hash_status"] = "completed"
hash_status = "completed"
else:
self._logger.error(f"Failed to calculate hash for {file_path}")
failures.append({"name": model.get("model_name", file_path or "Unknown"), "error": "Failed to calculate hash"})
processed += 1
handled_count += 1
continue
else:
self._logger.warning(f"Scanner does not support lazy hash calculation for {file_path}")
skipped_count += 1
processed += 1
handled_count += 1
continue
# Skip models without valid hash
if not model.get("sha256"):
self._logger.warning(f"Skipping model without hash: {file_path}")
skipped_count += 1
processed += 1
handled_count += 1
continue
await MetadataManager.hydrate_model_data(model)
result, _ = await self._metadata_sync.fetch_and_update_model(
# hydrate_model_data replaces model with .metadata.json content,
# which may lack sha256. Restore from cache and persist the fix.
if not model.get("sha256"):
model["sha256"] = sha256
model["hash_status"] = model.get("hash_status", hash_status)
data_to_save = model.copy()
data_to_save.pop("folder", None)
await MetadataManager.save_metadata(file_path, data_to_save)
result, error_msg = await self._metadata_sync.fetch_and_update_model(
sha256=model["sha256"],
file_path=model["file_path"],
model_data=model,
update_cache_func=self._service.scanner.update_single_model_cache,
)
if not result and error_msg and "Rate limited" in error_msg:
consecutive_rate_limits += 1
else:
consecutive_rate_limits = 0
if not result:
current_name = model.get("model_name", file_path or "Unknown")
failures.append({"name": current_name, "error": error_msg or "Unknown error"})
self._logger.warning("Failed to fetch metadata for %s: %s", current_name, error_msg)
if consecutive_rate_limits >= RATE_LIMIT_ABORT_THRESHOLD:
# The current model was attempted and failed due to rate limiting;
# count it before aborting so the summary is consistent.
processed += 1
handled_count += 1
self._logger.warning(
"Bulk metadata refresh aborted: %d consecutive rate limits detected. "
"Processed %d/%d models.",
consecutive_rate_limits,
processed,
total_to_process,
)
await emit(
"rate_limited",
)
return {
"success": False,
"message": f"Rate limit detected; {total_to_process - processed} models skipped",
"processed": processed,
"updated": success,
"total": total_models,
"failures": failures,
"failure_count": len(failures),
"skipped_count": skipped_count,
"elapsed_seconds": int(time.monotonic() - start_time),
}
if result:
success += 1
if original_name != model.get("model_name"):
needs_resort = True
processed += 1
handled_count += 1
await emit(
"processing",
processed=processed,
@@ -134,6 +217,9 @@ class BulkMetadataRefreshUseCase:
)
except Exception as exc: # pragma: no cover - logging path
processed += 1
handled_count += 1
current_name = model.get("model_name", model.get("file_path", "Unknown"))
failures.append({"name": current_name, "error": str(exc)})
self._logger.error(
"Error fetching CivitAI data for %s: %s",
model.get("file_path"),
@@ -150,7 +236,7 @@ class BulkMetadataRefreshUseCase:
f"{success} of {processed} processed {self._service.model_type}s (total: {total_models})"
)
return {"success": True, "message": message, "processed": processed, "updated": success, "total": total_models}
return {"success": True, "message": message, "processed": processed, "updated": success, "total": total_models, "failures": failures, "failure_count": len(failures), "skipped_count": skipped_count, "elapsed_seconds": int(time.monotonic() - start_time)}
@staticmethod
def _is_in_skip_path(folder: str, skip_paths: List[str]) -> bool:
+36 -3
View File
@@ -19,7 +19,7 @@ logger = logging.getLogger(__name__)
_WILDCARD_PATTERN = re.compile(r"__([\w\s.\-+/*\\]+?)__")
_OPTION_PATTERN = re.compile(r"{([^{}]*?)}")
_TRIGGER_WORD_PATTERN = re.compile(r"^trigger_words\d+$")
_WEIGHTED_OPTION_PATTERN = re.compile(r"^\s*([0-9.]+)::")
_WEIGHTED_OPTION_PATTERN = re.compile(r"^\s*-?\d+(\.\d+)?::")
_NUMERIC_PATTERN = re.compile(r"^-?\d+(\.\d+)?$")
@@ -390,7 +390,7 @@ class WildcardService:
) -> str | None:
keyword = _normalize_wildcard_key(raw_key)
if keyword in wildcard_dict:
return rng.choice(wildcard_dict[keyword])
return self._pick_weighted_or_plain(wildcard_dict[keyword], rng)
if "*" in keyword:
regex_pattern = keyword.replace("*", ".*").replace("+", r"\+")
@@ -400,7 +400,7 @@ class WildcardService:
if compiled.match(key):
aggregated.extend(values)
if aggregated:
return rng.choice(aggregated)
return self._pick_weighted_or_plain(aggregated, rng)
if "/" not in keyword:
fallback_keyword = _normalize_wildcard_key(f"*/{keyword}")
@@ -409,6 +409,39 @@ class WildcardService:
return None
def _pick_weighted_or_plain(
self, values: list[str], rng: random.Random
) -> str:
"""Pick a value from the list, respecting N::weight prefix if present.
When any value in the list uses the ``N::value`` weighted syntax with a
weight different from 1, the pick uses weighted random selection. When
no such weighting is present, a plain ``rng.choice`` is used (preserving
backward compatibility for unweighted wildcard files).
In either case the ``N::`` prefix is always stripped from the returned
value, matching the behaviour of ``{...}`` option groups.
"""
# Fast path: skip weighting logic entirely when no :: syntax exists
if not any("::" in v for v in values):
return rng.choice(values)
weighted_options: list[tuple[float, str]] = []
for value in values:
weight = 1.0
parts = value.split("::", 1)
if len(parts) == 2 and _is_numeric_string(parts[0].strip()):
weight = float(parts[0].strip())
weighted_options.append((weight, value))
any_weighted = any(w != 1.0 for w, _ in weighted_options)
if any_weighted:
picked = self._weighted_choice(weighted_options, rng)
else:
picked = rng.choice(values)
return self._strip_weight_prefix(picked)
def is_trigger_words_input(name: str) -> bool:
return bool(_TRIGGER_WORD_PATTERN.match(name))
+33 -1
View File
@@ -12,6 +12,7 @@ NODE_TYPES = {
"Lora Loader (LoraManager)": 1,
"Lora Stacker (LoraManager)": 2,
"WanVideo Lora Select (LoraManager)": 3,
"Create Hook LoRA (LoraManager)": 4,
}
# Default ComfyUI node color when bgcolor is null
@@ -31,6 +32,8 @@ PREVIEW_EXTENSIONS = [
".mp4",
".gif",
".webm",
".avif",
".jxl",
]
# Card preview image width
@@ -41,10 +44,24 @@ EXAMPLE_IMAGE_WIDTH = 832
# Supported media extensions for example downloads
SUPPORTED_MEDIA_EXTENSIONS = {
"images": [".jpg", ".jpeg", ".png", ".webp", ".gif"],
"images": [".jpg", ".jpeg", ".png", ".webp", ".gif", ".avif", ".jxl"],
"videos": [".mp4", ".webm"],
}
# Model weight file extensions recognised by scanners.
# This is the union of all scanner extensions (lora, checkpoint, embedding).
MODEL_FILE_EXTENSIONS = {
".safetensors",
".ckpt",
".pt",
".pt2",
".bin",
".pth",
".pkl",
".sft",
".gguf",
}
# Valid sub-types for each scanner type
VALID_LORA_SUB_TYPES = ["lora", "locon", "dora"]
VALID_CHECKPOINT_SUB_TYPES = ["checkpoint", "diffusion_model"]
@@ -145,6 +162,8 @@ DIFFUSION_MODEL_BASE_MODELS = frozenset(
"Qwen",
"ZImageBase",
"ZImageTurbo",
# Krea 2 — loaded via UNETLoader in ComfyUI
"Krea 2",
]
)
@@ -208,8 +227,21 @@ SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS = frozenset(
"Wan Video 2.5 I2V",
"Hunyuan Video",
"Anima",
"ACE Audio",
"Boogu",
"Ernie",
"Ernie Turbo",
"Grok",
"HappyHorse",
"HiDream-O1",
"Ideogram 4.0",
"Krea 2",
"Lens",
"MAI",
"Nucleus",
"Qwen 2",
"Upscaler",
"Wan Image 2.7",
"Wan Video 2.7",
]
)
+81 -24
View File
@@ -72,6 +72,7 @@ class _DownloadProgress(dict):
refreshed_models=set(),
failed_models=set(),
reprocessed_models=set(),
rate_limited_models=set(),
)
def snapshot(self) -> dict:
@@ -82,6 +83,7 @@ class _DownloadProgress(dict):
snapshot["refreshed_models"] = list(self["refreshed_models"])
snapshot["failed_models"] = list(self["failed_models"])
snapshot["reprocessed_models"] = list(self.get("reprocessed_models", set()))
snapshot["rate_limited_models"] = list(self.get("rate_limited_models", set()))
return snapshot
@@ -153,13 +155,15 @@ class DownloadManager:
# Step 3: Load progress file (I/O operation, done outside lock)
processed_models = set()
failed_models = set()
rate_limited_models = set()
try:
progress_file, processed_models, failed_models = await self._load_progress_file(output_dir)
progress_file, processed_models, failed_models, rate_limited_models = await self._load_progress_file(output_dir)
logger.debug(
"Loaded previous progress, %s models already processed, %s models marked as failed",
"Loaded previous progress, %s models already processed, %s models marked as failed, %s models rate-limited",
len(processed_models),
len(failed_models),
len(rate_limited_models),
)
except Exception as e:
logger.error(f"Failed to load progress file: {e}")
@@ -175,6 +179,7 @@ class DownloadManager:
self._progress.reset()
self._progress["processed_models"] = processed_models
self._progress["failed_models"] = failed_models
self._progress["rate_limited_models"] = rate_limited_models
self._stop_requested = False
self._progress["status"] = "running"
self._progress["start_time"] = time.time()
@@ -242,8 +247,8 @@ class DownloadManager:
"status": self._progress.snapshot(),
}
async def _load_progress_file(self, output_dir: str) -> tuple[str, set, set]:
"""Load progress file from disk. Returns (progress_file_path, processed_models, failed_models).
async def _load_progress_file(self, output_dir: str) -> tuple[str, set, set, set]:
"""Load progress file from disk. Returns (progress_file_path, processed_models, failed_models, rate_limited_models).
This is a separate async method to allow running in executor to avoid blocking event loop.
"""
@@ -252,8 +257,12 @@ class DownloadManager:
None, self._load_progress_file_sync, output_dir
)
def _load_progress_file_sync(self, output_dir: str) -> tuple[str, set, set]:
"""Synchronous implementation of progress file loading."""
def _load_progress_file_sync(self, output_dir: str) -> tuple[str, set, set, set]:
"""Synchronous implementation of progress file loading.
Returns:
tuple: (progress_file_path, processed_models, failed_models, rate_limited_models)
"""
progress_file = os.path.join(output_dir, ".download_progress.json")
progress_source = progress_file
@@ -289,6 +298,7 @@ class DownloadManager:
processed_models = set()
failed_models = set()
rate_limited_models = set()
if os.path.exists(progress_source):
try:
@@ -296,11 +306,11 @@ class DownloadManager:
saved_progress = json.load(f)
processed_models = set(saved_progress.get("processed_models", []))
failed_models = set(saved_progress.get("failed_models", []))
rate_limited_models = set(saved_progress.get("rate_limited_models", []))
except Exception:
# Return empty sets on error
pass
return progress_file, processed_models, failed_models
return progress_file, processed_models, failed_models, rate_limited_models
def _load_progress_sets_sync(self, progress_file: str) -> tuple[set, set]:
"""Load only the processed and failed model sets from progress file.
@@ -732,11 +742,13 @@ class DownloadManager:
success,
is_stale,
failed_images,
rate_limited_images,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash, model_name, images, model_dir, optimize, downloader
)
failed_urls: Set[str] = set(failed_images)
rate_limited_urls: Set[str] = set(rate_limited_images)
# If metadata is stale, try to refresh it
if is_stale and model_hash not in self._progress["refreshed_models"]:
@@ -760,6 +772,7 @@ class DownloadManager:
success,
_,
additional_failed,
additional_rate_limited,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
@@ -770,29 +783,50 @@ class DownloadManager:
)
failed_urls.update(additional_failed)
rate_limited_urls.update(additional_rate_limited)
self._progress["refreshed_models"].add(model_hash)
if failed_urls:
# Separate permanent failures from rate-limited ones
permanent_failures = failed_urls - rate_limited_urls
if permanent_failures:
await self._remove_failed_images_from_metadata(
model_hash,
model_name,
model_dir,
failed_urls,
permanent_failures,
scanner,
)
if failed_urls:
if rate_limited_urls:
self._progress["rate_limited_models"].add(model_hash)
logger.warning(
"%d example images for %s are rate-limited (429), will retry next time",
len(rate_limited_urls),
model_name,
)
# Clear failed_models so non-force runs can retry
if force and model_hash in self._progress["failed_models"]:
self._progress["failed_models"].discard(model_hash)
logger.info(
f"Removed {model_name} from failed_models after force retry with rate-limited images"
)
if rate_limited_urls:
# Don't mark as failed or fully processed — rate-limited
# images will be retried next time.
pass
elif permanent_failures:
self._progress["failed_models"].add(model_hash)
self._progress["processed_models"].add(model_hash)
logger.info(
"Removed %s failed example images for %s",
len(failed_urls),
len(permanent_failures),
model_name,
)
elif success:
self._progress["processed_models"].add(model_hash)
# Remove from failed_models if force mode enabled and model was previously failed
if force and model_hash in self._progress["failed_models"]:
self._progress["failed_models"].discard(model_hash)
logger.info(
@@ -850,6 +884,7 @@ class DownloadManager:
"processed_models": list(self._progress["processed_models"]),
"refreshed_models": list(self._progress["refreshed_models"]),
"failed_models": list(self._progress["failed_models"]),
"rate_limited_models": list(self._progress.get("rate_limited_models", set())),
"completed": self._progress["completed"],
"total": self._progress["total"],
"last_update": time.time(),
@@ -1155,11 +1190,13 @@ class DownloadManager:
success,
is_stale,
failed_images,
rate_limited_images,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash, model_name, images, model_dir, optimize, downloader
)
failed_urls: Set[str] = set(failed_images)
rate_limited_urls: Set[str] = set(rate_limited_images)
# If metadata is stale, try to refresh it
if is_stale and model_hash not in self._progress["refreshed_models"]:
@@ -1183,6 +1220,7 @@ class DownloadManager:
success,
_,
additional_failed_images,
additional_rate_limited,
) = await ExampleImagesProcessor.download_model_images_with_tracking(
model_hash,
model_name,
@@ -1192,21 +1230,35 @@ class DownloadManager:
downloader,
)
# Combine failed images from both attempts
failed_urls.update(additional_failed_images)
rate_limited_urls.update(additional_rate_limited)
self._progress["refreshed_models"].add(model_hash)
# For forced downloads, remove failed images from metadata
if failed_urls:
# Separate permanent failures from rate-limited ones
permanent_failures = failed_urls - rate_limited_urls
# Only remove permanently failed images from metadata
if permanent_failures:
await self._remove_failed_images_from_metadata(
model_hash, model_name, model_dir, failed_urls, scanner
model_hash, model_name, model_dir, permanent_failures, scanner
)
# Mark as processed
if (
success or failed_urls
): # Mark as processed if we successfully downloaded some images or removed failed ones
if rate_limited_urls:
self._progress["rate_limited_models"].add(model_hash)
logger.warning(
"%d example images for %s are rate-limited (429), will retry next time",
len(rate_limited_urls),
model_name,
)
# Mark as processed only when no rate-limited images remain
if rate_limited_urls:
pass
elif permanent_failures:
self._progress["processed_models"].add(model_hash)
self._progress["failed_models"].add(model_hash)
elif success:
self._progress["processed_models"].add(model_hash)
return True # Return True to indicate a remote download happened
@@ -1229,15 +1281,20 @@ class DownloadManager:
model_dir: str,
failed_images: Iterable[str],
scanner,
error_type: str = "not_found",
) -> None:
"""Mark failed images in model metadata so they won't be retried."""
"""Mark failed images in model metadata so they won't be retried.
Args:
error_type: Reason string stored in the image's ``downloadError`` field
(default ``"not_found"``).
"""
failed_set: Set[str] = {url for url in failed_images if url}
if not failed_set:
return
try:
# Get current model data
model_data = await MetadataUpdater.get_updated_model(model_hash, scanner)
if not model_data:
logger.warning(
@@ -1268,7 +1325,7 @@ class DownloadManager:
continue
image["downloadFailed"] = True
image.setdefault("downloadError", "not_found")
image.setdefault("downloadError", error_type)
logger.debug(
"Marked example image %s for %s as failed due to missing remote asset",
image_url,
+17 -33
View File
@@ -475,13 +475,19 @@ class MetadataUpdater:
return False
model_folder = get_model_folder(model_hash)
if not model_folder:
if not model_folder or not os.path.isdir(model_folder):
return False
civitai = getattr(metadata, "civitai", None)
if not isinstance(civitai, dict):
return False
# Read the directory listing once so every image entry reuses it.
try:
dir_entries = os.listdir(model_folder)
except OSError:
dir_entries = []
has_changes = False
custom_images = civitai.get("customImages")
@@ -493,24 +499,15 @@ class MetadataUpdater:
if not img_id:
continue
if not os.path.isdir(model_folder):
prefix = f"custom_{img_id}"
found = any(
f.startswith(prefix) and os.path.isfile(
os.path.join(model_folder, f)
)
for f in dir_entries
)
if not found:
stale.append(idx)
else:
found = False
try:
prefix = f"custom_{img_id}"
for fname in os.listdir(model_folder):
if fname.startswith(prefix) and os.path.isfile(
os.path.join(model_folder, fname)
):
found = True
break
except OSError:
stale.append(idx)
continue
if not found:
stale.append(idx)
if stale:
for idx in reversed(stale):
@@ -532,22 +529,9 @@ class MetadataUpdater:
# is gone.
continue
if not os.path.isdir(model_folder):
prefix = f"image_{idx}."
if not any(f.startswith(prefix) for f in dir_entries):
stale.append(idx)
else:
found = False
try:
prefix = f"image_{idx}."
for fname in os.listdir(model_folder):
if fname.startswith(prefix):
found = True
break
except OSError:
stale.append(idx)
continue
if not found:
stale.append(idx)
if stale:
for idx in reversed(stale):
+96 -1
View File
@@ -3,9 +3,16 @@ import logging
import os
import re
import json
import shutil
from ..services.settings_manager import get_settings_manager
from ..services.service_registry import ServiceRegistry
from ..utils.example_images_paths import iter_library_roots
from ..utils.example_images_paths import (
get_example_images_root,
is_hash_folder,
iter_library_roots,
uses_library_scoped_folders,
_library_folder_has_only_hash_dirs,
)
from ..utils.metadata_manager import MetadataManager
from ..utils.example_images_processor import ExampleImagesProcessor
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
@@ -36,6 +43,90 @@ settings = _SettingsProxy()
class ExampleImagesMigration:
"""Handles migrations for example images naming conventions"""
@staticmethod
def _consolidate_library_folders():
"""Move hash folders from library-named subdirectories back to root.
When a user switches from multi-library mode back to single-library
mode, example images previously stored under e.g.
``<root>/default/<hash>/`` need to be moved back to
``<root>/<hash>/``. Running this once at startup removes the need
for ``get_model_folder()`` to perform directory scans on every
request.
"""
if uses_library_scoped_folders():
return
root = get_example_images_root()
if not root or not os.path.isdir(root):
return
moved: list[str] = []
cleaned: list[str] = []
try:
for entry in os.listdir(root):
# Fast regex checks first — no filesystem I/O.
if is_hash_folder(entry) or entry == "_deleted":
continue
entry_path = os.path.join(root, entry)
if not os.path.isdir(entry_path):
continue
if not _library_folder_has_only_hash_dirs(entry_path):
continue
try:
for hash_entry in os.listdir(entry_path):
hash_path = os.path.join(entry_path, hash_entry)
if not os.path.isdir(hash_path) or not is_hash_folder(hash_entry):
continue
target = os.path.join(root, hash_entry)
if not os.path.exists(target):
try:
shutil.move(hash_path, target)
moved.append(hash_entry)
except (OSError, shutil.Error) as exc:
logger.error(
"Failed to move '%s''%s': %s",
hash_path, target, exc,
)
except OSError as exc:
logger.error(
"Failed to list library subdirectory '%s': %s",
entry_path, exc,
)
try:
remaining = os.listdir(entry_path)
except OSError:
remaining = []
if not remaining:
try:
os.rmdir(entry_path)
cleaned.append(entry)
except OSError as exc:
logger.debug(
"Could not remove empty library dir '%s': %s",
entry_path, exc,
)
except OSError as exc:
logger.error(
"Failed to list example images root during consolidation: %s",
exc,
)
if moved:
logger.info(
"Consolidated %d example image folder(s) to root",
len(moved),
)
if cleaned:
logger.info(
"Removed %d empty library directories",
len(cleaned),
)
@staticmethod
async def check_and_run_migrations():
"""Check if migrations are needed and run them in background"""
@@ -44,6 +135,10 @@ class ExampleImagesMigration:
logger.debug("No example images path configured or path doesn't exist, skipping migrations")
return
# Run library-to-root consolidation once at startup so the hot
# path (get_model_folder) stays a pure-path computation.
ExampleImagesMigration._consolidate_library_folders()
for library_name, library_path in iter_library_roots():
if not library_path or not os.path.exists(library_path):
continue
+76 -1
View File
@@ -12,6 +12,18 @@ from ..services.settings_manager import get_settings_manager
_HEX_PATTERN = re.compile(r"[a-fA-F0-9]{64}")
# Filesystem/metadata files that are never created by the example images system
# and are safe to ignore during validation. The cleanup service only operates on
# directories, so these files pose no data-loss risk.
_SAFE_FILENAMES: frozenset[str] = frozenset({
".DS_Store", # macOS folder metadata
"Thumbs.db", # Windows thumbnail cache
"desktop.ini", # Windows folder customization
".localized", # macOS folder name localization
".gitkeep", # Placeholder to keep empty dirs in git
".gitignore", # Git ignore rules
})
logger = logging.getLogger(__name__)
@@ -71,7 +83,12 @@ def ensure_library_root_exists(library_name: Optional[str] = None) -> str:
def get_model_folder(model_hash: str, library_name: Optional[str] = None) -> str:
"""Return the folder path for a model's example images."""
"""Return the folder path for a model's example images.
Multi-library single-library consolidation is handled once at startup by
``ExampleImagesMigration._consolidate_library_folders`` this function is a
pure path computation on the hot path (no directory scans).
"""
if not model_hash:
return ""
@@ -180,6 +197,22 @@ def is_hash_folder(name: str) -> bool:
return bool(_HEX_PATTERN.fullmatch(name or ""))
def _is_safe_ignorable_entry(item: str, item_path: str) -> bool:
"""Return True if *item* is a harmless system/hidden file we can skip.
These files are never created by the example images system and are safe to
ignore because the cleanup/delete operations only act on **directories**,
never on individual files (other than ``.download_progress.json``).
"""
if item in _SAFE_FILENAMES:
return True
# Hide Unix hidden files (dotfiles) that are regular files,
# since the cleanup system never deletes or moves files.
if item.startswith(".") and os.path.isfile(item_path):
return True
return False
def is_valid_example_images_root(folder_path: str) -> bool:
"""Check whether a folder looks like a dedicated example images root."""
@@ -190,9 +223,16 @@ def is_valid_example_images_root(folder_path: str) -> bool:
for item in items:
item_path = os.path.join(folder_path, item)
# .download_progress.json is an expected metadata file — check before
# the generic dotfile rule so it stays explicitly documented.
if item == ".download_progress.json" and os.path.isfile(item_path):
continue
# Skip harmless system/hidden files — cleanup only touches directories
if _is_safe_ignorable_entry(item, item_path):
continue
if os.path.isdir(item_path):
if is_hash_folder(item):
continue
@@ -211,6 +251,41 @@ def is_valid_example_images_root(folder_path: str) -> bool:
return True
def find_non_compliant_items_in_example_images_root(folder_path: str) -> list[str]:
"""Return the names of items that prevent *folder_path* from being a valid
example images root, or an empty list if the folder is valid.
This mirrors ``is_valid_example_images_root`` but **returns** the offending
names instead of a boolean, so callers can produce actionable error messages.
"""
try:
items = os.listdir(folder_path)
except OSError as exc:
return [f"<cannot list directory: {exc}>"]
offending: list[str] = []
for item in items:
item_path = os.path.join(folder_path, item)
# Same skip rules as is_valid_example_images_root
if item == ".download_progress.json" and os.path.isfile(item_path):
continue
if _is_safe_ignorable_entry(item, item_path):
continue
if os.path.isdir(item_path):
if is_hash_folder(item):
continue
if item == "_deleted":
continue
if _library_folder_has_only_hash_dirs(item_path):
continue
offending.append(item)
return offending
def _library_folder_has_only_hash_dirs(path: str) -> bool:
"""Return True when a library subfolder only contains hash folders or metadata files."""
+98 -39
View File
@@ -1,3 +1,4 @@
import asyncio
import logging
import os
import re
@@ -62,6 +63,10 @@ class ExampleImagesProcessor:
return '.gif'
elif content.startswith(b'RIFF') and b'WEBP' in content[:12]:
return '.webp'
elif len(content) >= 12 and content[4:8] == b'ftyp' and b'avif' in content[8:24]:
return '.avif'
elif content.startswith(b'\x00\x00\x00\x0cJXL \x0d\x0a\x87\x0a'):
return '.jxl'
elif content.startswith(b'\x00\x00\x00\x18ftypmp4') or content.startswith(b'\x00\x00\x00\x20ftypmp4'):
return '.mp4'
elif content.startswith(b'\x1A\x45\xDF\xA3'):
@@ -75,6 +80,8 @@ class ExampleImagesProcessor:
'image/png': '.png',
'image/gif': '.gif',
'image/webp': '.webp',
'image/avif': '.avif',
'image/jxl': '.jxl',
'video/mp4': '.mp4',
'video/webm': '.webm',
'video/quicktime': '.mov'
@@ -188,16 +195,22 @@ class ExampleImagesProcessor:
return model_success, False # (success, is_metadata_stale)
@staticmethod
def _extract_retry_after(error_message: str) -> int:
if not error_message:
return 60
match = re.search(r"retry after (\d+)s", str(error_message))
if match:
return max(1, int(match.group(1)))
return 60
@staticmethod
async def download_model_images_with_tracking(model_hash, model_name, model_images, model_dir, optimize, downloader):
"""Download images for a single model with tracking of failed image URLs
Returns:
tuple: (success, is_stale_metadata, failed_images) - whether download was successful, whether metadata is stale, list of failed image URLs
"""
model_success = True
failed_images = []
rate_limited_images = []
any_successful_download = False
for i, image in enumerate(model_images):
image_url = image.get('url')
if not image_url:
@@ -215,64 +228,110 @@ class ExampleImagesProcessor:
original_url = image_url
if optimize and 'civitai.com' in image_url:
image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url)
# Download the file first to determine the actual file type
try:
logger.debug(f"Downloading media file {i} for {model_name}")
# Download using the unified downloader with headers
success, content, headers = await downloader.download_to_memory(
async def _attempt_download() -> tuple:
logger.debug("Downloading media file %s for %s", i, model_name)
return await downloader.download_to_memory(
image_url,
use_auth=False, # Example images don't need auth
return_headers=True
use_auth=False,
return_headers=True,
)
try:
success, content, headers = await _attempt_download()
if success:
# Determine file extension from content or headers
media_ext = ExampleImagesProcessor._get_file_extension_from_content_or_headers(
content, headers, original_url, image.get("type")
)
# Check if the detected file type is supported
is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS['images']
is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
if not (is_image or is_video):
logger.debug(f"Skipping unsupported file type: {media_ext}")
logger.debug("Skipping unsupported file type: %s", media_ext)
continue
# Use 0-based indexing with the detected extension
save_filename = f"image_{i}{media_ext}"
save_path = os.path.join(model_dir, save_filename)
# Check if already downloaded
if os.path.exists(save_path):
logger.debug(f"File already exists: {save_path}")
logger.debug("File already exists: %s", save_path)
continue
# Save the file
with open(save_path, 'wb') as f:
f.write(content)
any_successful_download = True
elif ExampleImagesProcessor._is_not_found_error(content):
error_msg = f"Failed to download file: {image_url}, status code: 404 - Model metadata might be stale"
logger.warning(error_msg)
model_success = False # Mark the model as failed due to 404 error
failed_images.append(image_url) # Track failed URL
# Return early to trigger metadata refresh attempt
return False, True, failed_images # (success, is_metadata_stale, failed_images)
model_success = False
failed_images.append(image_url)
return False, True, failed_images, rate_limited_images
elif "Rate limited (429)" in str(content):
max_attempts = 3
for attempt in range(1, max_attempts + 1):
wait = ExampleImagesProcessor._extract_retry_after(str(content)) * (2 ** (attempt - 1))
logger.warning(
"Rate limited (429) for %s, retry %d/%d after %ds",
image_url, attempt, max_attempts, wait,
)
await asyncio.sleep(wait)
success, content, headers = await _attempt_download()
if success:
media_ext = ExampleImagesProcessor._get_file_extension_from_content_or_headers(
content, headers, original_url, image.get("type")
)
is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS['images']
is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
if not (is_image or is_video):
logger.debug("Skipping unsupported file type: %s", media_ext)
break
save_filename = f"image_{i}{media_ext}"
save_path = os.path.join(model_dir, save_filename)
if os.path.exists(save_path):
logger.debug("File already exists: %s", save_path)
break
with open(save_path, 'wb') as f:
f.write(content)
any_successful_download = True
break
elif "Rate limited (429)" in str(content):
continue
elif ExampleImagesProcessor._is_not_found_error(content):
logger.warning("Failed to download file: %s, status code: 404", image_url)
model_success = False
failed_images.append(image_url)
break
else:
logger.warning("Failed to download file: %s, error: %s", image_url, content)
model_success = False
failed_images.append(image_url)
break
else:
logger.warning(
"Giving up on %s after %d retries due to rate limiting",
image_url, max_attempts,
)
rate_limited_images.append(image_url)
model_success = False
else:
error_msg = f"Failed to download file: {image_url}, error: {content}"
logger.warning(error_msg)
model_success = False # Mark the model as failed
failed_images.append(image_url) # Track failed URL
model_success = False
failed_images.append(image_url)
except Exception as e:
error_msg = f"Error downloading file {image_url}: {str(e)}"
logger.error(error_msg)
model_success = False # Mark the model as failed
failed_images.append(image_url) # Track failed URL
return model_success, False, failed_images # (success, is_metadata_stale, failed_images)
model_success = False
failed_images.append(image_url)
return any_successful_download or model_success, False, failed_images, rate_limited_images
@staticmethod
async def process_local_examples(model_file_path, model_file_name, model_name, model_dir, optimize):
+117 -7
View File
@@ -1,17 +1,125 @@
import json
import logging
import os
import struct
from io import BytesIO
from typing import Any, Optional
import piexif
from PIL import Image, PngImagePlugin
try:
import brotli
_BROTLI_AVAILABLE = True
except ImportError:
brotli = None
_BROTLI_AVAILABLE = False
logger = logging.getLogger(__name__)
class ExifUtils:
"""Utility functions for working with EXIF data in images"""
@staticmethod
def _parse_isobmff_boxes(data: bytes, offset: int = 0) -> list[dict]:
boxes = []
while offset + 8 <= len(data):
size = struct.unpack('>I', data[offset:offset + 4])[0]
box_type = data[offset + 4:offset + 8]
if size == 0:
break
if size < 8 or offset + size > len(data):
break
box_data = data[offset + 8:offset + size]
boxes.append({'type': box_type, 'data': box_data, 'size': size})
offset += size
return boxes
@staticmethod
def _is_jxl_container(data: bytes) -> bool:
if len(data) < 32:
return False
return (
struct.unpack('>I', data[:4])[0] == 12
and data[4:8] == b'JXL '
and data[8:12] == bytes([0x0d, 0x0a, 0x87, 0x0a])
and struct.unpack('>I', data[12:16])[0] >= 16
and data[16:20] == b'ftyp'
and data[20:24] == b'jxl '
)
@staticmethod
def _is_avif_container(data: bytes) -> bool:
if len(data) < 16:
return False
for box in ExifUtils._parse_isobmff_boxes(data):
if box['type'] == b'ftyp' and b'avif' in box['data']:
return True
return False
# Max decompressed size for brotli metadata (2 MB)
_BROTLI_MAX_DECOMPRESSED = 2 * 1024 * 1024
@staticmethod
def _extract_isobmff_brotli(image_path: str) -> Optional[dict]:
try:
with open(image_path, 'rb') as f:
data = f.read()
except Exception:
return None
if ExifUtils._is_jxl_container(data):
boxes = ExifUtils._parse_isobmff_boxes(data, offset=12)
elif ExifUtils._is_avif_container(data):
boxes = ExifUtils._parse_isobmff_boxes(data)
else:
return None
brob = None
for box in boxes:
if box['type'] == b'brob':
brob = box
break
if brob is None:
return None
payload = brob['data']
if payload[:4] != b'comf':
return None
compressed = payload[4:]
if _BROTLI_AVAILABLE:
try:
decompressed = brotli.decompress(compressed)
if len(decompressed) > ExifUtils._BROTLI_MAX_DECOMPRESSED:
logger.warning(
"Brotli metadata too large (%d bytes, max %d), ignoring",
len(decompressed),
ExifUtils._BROTLI_MAX_DECOMPRESSED,
)
decompressed = None
except Exception:
decompressed = None
else:
decompressed = None
raw = decompressed if decompressed is not None else compressed
try:
meta = json.loads(raw.decode('utf-8'))
except Exception:
return None
result = {"parameters": None, "prompt": None, "workflow": None, "comment": None}
if isinstance(meta.get("prompt"), (dict, list)):
result["prompt"] = json.dumps(meta["prompt"])
elif isinstance(meta.get("prompt"), str):
result["prompt"] = meta["prompt"]
if isinstance(meta.get("workflow"), (dict, list)):
result["workflow"] = json.dumps(meta["workflow"])
elif isinstance(meta.get("workflow"), str):
result["workflow"] = meta["workflow"]
return result
@staticmethod
def _decode_user_comment(user_comment: Any) -> Optional[str]:
if user_comment is None:
@@ -43,6 +151,12 @@ class ExifUtils:
"comment": None,
}
ext = os.path.splitext(image_path)[1].lower()
if ext in ('.avif', '.jxl'):
brotli_meta = ExifUtils._extract_isobmff_brotli(image_path)
if brotli_meta:
return brotli_meta
with Image.open(image_path) as img:
info = getattr(img, "info", {}) or {}
@@ -149,7 +263,6 @@ class ExifUtils:
Optional[str]: Extracted metadata or None if not found
"""
try:
# Skip for video files
if image_path:
ext = os.path.splitext(image_path)[1].lower()
if ext in ['.mp4', '.webm']:
@@ -177,10 +290,9 @@ class ExifUtils:
str: Path to the updated image
"""
try:
# Skip for video files
if image_path:
ext = os.path.splitext(image_path)[1].lower()
if ext in ['.mp4', '.webm']:
if ext in ['.mp4', '.webm', '.avif', '.jxl']:
return image_path
metadata_fields = ExifUtils._load_structured_metadata(image_path)
@@ -212,10 +324,9 @@ class ExifUtils:
def append_recipe_metadata(image_path, recipe_data) -> str:
"""Append recipe metadata to an image's EXIF data"""
try:
# Skip for video files
if image_path:
ext = os.path.splitext(image_path)[1].lower()
if ext in ['.mp4', '.webm']:
if ext in ['.mp4', '.webm', '.avif', '.jxl']:
return image_path
# First, extract existing metadata
@@ -327,10 +438,9 @@ class ExifUtils:
Tuple of (optimized_image_data, extension)
"""
try:
# Skip for video files early if it's a file path
if isinstance(image_data, str) and os.path.exists(image_data):
ext = os.path.splitext(image_data)[1].lower()
if ext in ['.mp4', '.webm']:
if ext in ['.mp4', '.webm', '.avif', '.jxl']:
try:
with open(image_data, 'rb') as f:
return f.read(), ext
+15 -1
View File
@@ -34,12 +34,26 @@ def _get_hash_chunk_size_bytes() -> int:
async def calculate_sha256(file_path: str) -> str:
"""Calculate SHA256 hash of a file (full file content)."""
"""Calculate SHA256 hash of a file (full file content).
Uses ``posix_fadvise`` with ``POSIX_FADV_DONTNEED`` to avoid polluting the OS page
cache critical on WSL where cached file pages live inside the VM and are not
accounted for in guest ``used`` memory, causing VmmemWSL to balloon.
On Windows/macOS where ``posix_fadvise`` is not available the hint is silently
skipped.
"""
sha256_hash = hashlib.sha256()
chunk_size = _get_hash_chunk_size_bytes()
with open(file_path, "rb") as f:
fd = f.fileno()
for byte_block in iter(lambda: f.read(chunk_size), b""):
sha256_hash.update(byte_block)
# Evict pages after reading so the data doesn't linger in the kernel page
# cache — on WSL this otherwise appears as unreclaimable VmmemWSL growth.
# Guard against platforms (Windows, macOS) that lack posix_fadvise.
if hasattr(os, "posix_fadvise") and hasattr(os, "POSIX_FADV_DONTNEED"):
os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
return sha256_hash.hexdigest()

Some files were not shown because too many files have changed in this diff Show More