Compare commits

...

103 Commits

Author SHA1 Message Date
Will Miao 27027c4497 refactor(recipes): reuse shared local hash cache in create-from-example 2026-08-08 22:45:29 +08:00
Will Miao 86c85c08ec feat(recipes): pass local hash cache to remote and url recipe imports 2026-08-08 22:13:19 +08:00
Will Miao 196c8ffc3e feat(recipes): match civitai image hash sections against local hash cache 2026-08-08 22:12:47 +08:00
Will Miao cfc95ee02a feat(recipes): pass local hash cache through analysis recipe parsing 2026-08-08 22:11:50 +08:00
Will Miao 479fa36997 feat(recipes): add version-cached local hash cache builder 2026-08-08 22:04:09 +08:00
Will Miao 3e1216e9bc feat(recipes): add cache version counter to model scanners 2026-08-08 21:57:36 +08:00
Will Miao 007883b7d1 fix(recipes): backfill lora cache item by autov2/autov3 hash too 2026-08-08 21:51:34 +08:00
Will Miao dc9200a12c fix(recipes): match recipe-format lora cache item by autov2/autov3 hash 2026-08-08 21:50:33 +08:00
Will Miao d2f955266d fix(types): resolve pre-existing basedpyright errors in tests
Fix ~790 basedpyright errors across the test suite:
- Type stub subclasses of real production classes with super().__init__()
- Add missing generic type arguments and Dict[str, Any] annotations
- Add None guards before subscript/member access
- Adapt tests to production API changes (removed dead handlers,
  PersistentModelCache.get_default, _i18n_filter_added location)
2026-08-08 20:12:59 +08:00
Will Miao 8e724538bd fix(types): resolve pre-existing basedpyright errors in py/ and standalone.py
Fix ~950 basedpyright errors across the backend:
- Convert ineffective # type: ignore comments to # pyright: ignore[rule]
- Add missing generic type arguments (Dict[str, Any], list[Any], ...)
- Annotate dynamic dict literals and runtime-initialized attributes
- Widen CivitAI provider tuple signatures in recipe parsers
- Remove dead LoraRoutes handlers calling nonexistent LoraService methods
- Suppress unavoidable ServiceRegistry import cycles (basedpyright counts
  function-local imports as cycle edges)
2026-08-08 20:12:52 +08:00
Will Miao 6fcdeb799d feat(metadata): resolve AutoV3 at download time without waiting for backfill
- Read AutoV3 directly from the downloaded file's own file_info hashes
  (no SHA256 cross-matching against version_info.files, so the value is
  captured even when the API omits SHA256)
- Extract normalize_autov3() validation helper shared with the
  sha256-matching autov3_from_civitai_files path
- Fall back to the embedded safetensors header hash at download
  completion; mark '' (checked-unavailable) so the startup backfill
  query (autov3 IS NULL) never revisits the row
- Clear archive-level AutoV3 for zip-extracted models so per-file
  header resolution applies to every extracted model
2026-08-08 15:20:43 +08:00
Will Miao 97b9b1f62b feat(metadata): add CivitAI AutoV3 hash support across all storage layers
- Three-state autov3 field (not-checked / checked-unavailable / 12-hex value)
  in .metadata.json sidecars, in-memory ModelHashIndex, and SQLite
  (models.autov3 column + autov3_index table) with column-presence migration
- Background self-terminating backfill for legacy rows: per-model-type
  concurrency guard, executor-offloaded I/O, Civitai-first resolution
  (SHA256-matched version file) falling back to the embedded safetensors
  header hash
- Civitai-first propagation on metadata refresh, scan, and download paths;
  reject the empty-string SHA256 placeholder and strip OneTrainer 0x prefix
- List API hash filters and hash index lookups accept 12-char AutoV3
- Cap safetensors header reads at 64 MiB to prevent crafted-file allocation
- Prevent stale AutoV3 mappings on file replacement while preserving them on
  same-file re-registration (lazy-hash completion)
2026-08-08 14:30:34 +08:00
Will Miao 4bf9a4b640 refactor(download): rename locationStep id to downloadLocationStep
The download modal's step shared the 'locationStep' id with the import
modal, so getElementById('locationStep') could resolve to the wrong
element depending on template include order. The import flow relied on
an injected display:block !important rule to work around it.

Rename the download modal's step id and update all references so each
modal owns a unique step id.
2026-08-08 08:50:36 +08:00
Will Miao c5088772e8 fix(ui): pin import modal action buttons with sticky footer
Make the import modal a flex column with a scrollable step area so the
Back/Import buttons stay visible on short viewports (1080p / 150% zoom)
instead of being cut off at the bottom of the scroll flow.

Also reset step scroll positions via class since 'locationStep' has a
duplicate id in the download modal template.
2026-08-08 08:49:20 +08:00
Will Miao 56acefbd6c feat(autocomplete): search loras within active filters of LoRA Manager page
Add /af and /noaf toggle commands (plus /activefilters aliases) to the
loras autocomplete widget. When enabled (default off), suggestions are
matched within the active filters (folder, base model, tags, auto-tags,
license, tag logic) persisted by the LoRA Manager page in localStorage,
keeping the match pool consistent with the list endpoint, including the
global show_only_sfw setting.

Backend: /lm/{prefix}/relative-paths accepts the filter query params and
pre-filters the scanner cache with ModelFilterSet. The presence of the
recursive param signals the filter pipeline to run even without concrete
filters so global settings stay in parity with the list endpoint.
2026-08-07 20:07:34 +08:00
Will Miao 5ab06c4aae docs: rename "Standalone Web UI" to "LoRA Manager Web UI" in AGENTS.md 2026-08-07 17:57:01 +08:00
Will Miao c11f4b5c68 feat(ui): widen filter panel and preset name limit 2026-08-07 16:53:14 +08:00
Will Miao 86376284f4 fix(ui): clamp filter panel height to viewport 2026-08-07 16:53:04 +08:00
Will Miao 2b8a2fc7d8 feat(filters): remove preset count limit 2026-08-07 16:52:53 +08:00
Will Miao f26e1b41c8 fix(i18n): translate zh-TW api key placeholder 2026-08-07 16:25:40 +08:00
Will Miao c1671af99f feat(downloads): translate batch download summary strings 2026-08-07 16:23:55 +08:00
Will Miao ac7707d0f6 fix(cards): clear model-card min-width on the item element itself 2026-08-07 16:09:51 +08:00
Will Miao 381cd710a2 feat(recipes): translate recipes layout setting strings 2026-08-07 15:36:30 +08:00
Will Miao ad0d18cb79 chore: ignore .playwright-mcp working directory 2026-08-07 15:33:43 +08:00
Will Miao 7980ee77d0 perf(recipes): batch preview dimension reads via asyncio.gather 2026-08-07 15:31:13 +08:00
Will Miao 916b8bb327 fix(recipes): skip stale scroller re-enable on deferred layout switch 2026-08-07 14:03:13 +08:00
Will Miao 87e3d4dea9 feat(recipes): wire recipes layout switch event and rebuild 2026-08-07 12:49:30 +08:00
Will Miao 76a913f5e0 feat(recipes): complete MasonryScroller public API parity with VirtualScroller 2026-08-07 12:40:28 +08:00
Will Miao d8c192e647 feat(recipes): branch masonry scroller instantiation for recipes page 2026-08-07 12:38:17 +08:00
Will Miao c453437620 feat(recipes): add MasonryScroller with column-based virtual scrolling 2026-08-07 12:31:13 +08:00
Will Miao 720fa6d909 feat(recipes): expose preview width/height in recipe listing API 2026-08-07 11:59:26 +08:00
Will Miao b4f71089f4 feat(recipes): add recipes_layout setting (grid|masonry) with i18n 2026-08-07 11:49:12 +08:00
Will Miao 83e6657ead feat(recipes): add get_image_dimensions helper with LRU cache 2026-08-07 11:47:28 +08:00
Will Miao 7ea6df4111 feat(downloads): default to latest version when URL lacks modelVersionId
Auto-select the first (newest) version for URLs without an explicit
modelVersionId, matching the existing batch flow, so users can proceed
to location/download without manually picking a version.
2026-08-07 11:24:45 +08:00
Will Miao d9ab92602a feat(downloads): show failure summary modal for single downloads too 2026-08-07 10:49:14 +08:00
pixelpaws 5ffadaed31 Merge pull request #1054 from willmiao/feat/gemini-provider
feat(llm): add Gemini as a preset AI provider
2026-08-07 10:30:45 +08:00
Will Miao 24f5f7df5d feat(llm): add Gemini as a preset AI provider 2026-08-07 10:27:51 +08:00
Will Miao daf01fb1d6 feat(downloads): show batch download summary with failure details and retry 2026-08-07 10:23:17 +08:00
Will Miao 0f11b6def9 fix(recipes): allow recipes storage path on a different drive (Windows)
os.path.commonpath raises ValueError for paths on different Windows
drives. Treat that as no common root so cross-drive recipes migrations
succeed instead of failing with 'Invalid recipes path change'.
2026-08-06 22:18:24 +08:00
Will Miao 7df83f44b8 feat(SaveImageLM): add add_loras_to_prompt toggle to restore legacy lora syntax line in metadata 2026-08-06 15:58:18 +08:00
Will Miao 169fa7bed6 fix(vue-widgets): resolve pre-existing typecheck errors 2026-08-06 15:33:02 +08:00
Will Miao 027b504fe8 refactor(autocomplete): remove unused custom_words and embeddings modelTypes 2026-08-06 15:28:58 +08:00
Will Miao 186ef4da78 refactor(ui): group example image download actions into a submenu
Move the 'Download Missing' / 'Re-process All' example image actions
under a single 'Download Example Images' submenu item in the single-model
and bulk context menus, matching the existing send-to-workflow submenu
pattern. Shorten the submenu labels and update all locale translations.
2026-08-03 21:18:05 +08:00
pixelpaws dc674098e7 Merge pull request #1050 from willmiao/fix/recipes-bulk-content-rating
fix(recipes): enable bulk content rating for selected recipes
2026-08-03 20:58:24 +08:00
Will Miao 9087b4b07c feat(example-images): add missing-only download path and skip existing files
Split the single-model and bulk context menu actions into 'Download
Missing Example Images' (regular endpoint, skips already-processed
models) and 'Re-process Example Images' (force endpoint, retries
failed models).

- start_download accepts model_hashes so a selected subset can be
  processed with the progress-aware skip logic; explicitly targeted
  models bypass the failed/processed model-level guards so per-image
  gaps are filled
- pre-download existence check in the processor skips network requests
  for image files already on disk across all download paths
- force download retries previously failed models and clears their
  failed status on success
- add i18n keys for the new menu items across all locales
2026-08-03 20:52:46 +08:00
Will Miao 8e45c22d7a fix(recipes): enable bulk content rating for selected recipes 2026-08-03 19:31:58 +08:00
Will Miao 191c4e03cd feat(metadata-overwrite): support wired MODEL input on model field
The model field now accepts either a manual string or a MODEL connection.
When wired, the model name is extracted from the patcher's
cached_patcher_init (registered by core loaders load_checkpoint_guess_config
and load_diffusion_model, preserved through LoRA clones) and converted to a
ComfyUI-style relative name via config model roots.

- model input declared as "STRING,MODEL" with widgetType STRING, so the
  text widget and the dual-type connection slot coexist; non-STRING/MODEL
  links are rejected by frontend and backend type validation
- UNETLoaderLM GGUF branch now registers a custom cached_patcher_init reload
  factory so GGUF models participate in name extraction and ModelPatcher
  deepclone/dynamic machinery
- shared collect_overwrite_params() helper keeps the node and the metadata
  extractor conversion logic in sync; extraction failures are logged instead
  of silently dropping the overwrite
2026-08-03 16:44:03 +08:00
Will Miao ab4154c57d feat(ui): add seeded random sort option to model pages (#1049) 2026-08-03 15:02:49 +08:00
Will Miao 28e93d12ff fix(example-images): use in-place cache sync and bulk pending-check index for large libraries 2026-08-03 12:04:56 +08:00
Will Miao 75e63c758b feat(api): add cursor-based pagination to civitai user-models endpoint 2026-08-03 11:07:06 +08:00
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
316 changed files with 39178 additions and 22703 deletions
+1
View File
@@ -25,6 +25,7 @@ model_cache/
reasonix.toml
.reasonix/
.codegraph/
.playwright-mcp/
# Vue widgets development cache (but keep build output)
vue-widgets/node_modules/
+10 -4
View File
@@ -31,7 +31,7 @@ COVERAGE_FILE=coverage/backend/.coverage pytest \
--cov-report=xml:coverage/backend/coverage.xml
```
### Frontend Development (Standalone Web UI)
### Frontend Development (LoRA Manager Web UI)
```bash
npm install
@@ -137,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
@@ -148,9 +154,9 @@ npm run test:coverage # Generate coverage report
## Frontend UI Architecture
### 1. Standalone Web UI
### 1. LoRA Manager Web UI
- Location: `./static/` and `./templates/`
- Tech: Vanilla JS + CSS, served by standalone server
- Tech: Vanilla JS + CSS, served by the hosting server (ComfyUI app in plugin mode, `standalone.py` in standalone mode)
- Tests via npm in root directory
### 2. ComfyUI Custom Node Widgets
+2 -2
View File
File diff suppressed because one or more lines are too long
+10
View File
@@ -17,6 +17,8 @@ try: # pragma: no cover - import fallback for pytest collection
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
@@ -62,6 +64,12 @@ except (
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 = {
@@ -83,6 +91,8 @@ NODE_CLASS_MAPPINGS = {
LoraCyclerLM.NAME: LoraCyclerLM,
LoraInfoLM.NAME: LoraInfoLM,
LoraSyntaxToPath.NAME: LoraSyntaxToPath,
CreateHookLoraLM.NAME: CreateHookLoraLM,
MetadataOverwriteLM.NAME: MetadataOverwriteLM,
}
WEB_DIRECTORY = "./web/comfyui"
+313 -291
View File
File diff suppressed because it is too large Load Diff
+4
View File
@@ -39,6 +39,7 @@ These fields are present in all model metadata files.
| `metadata_source` | string\|null | ❌ No | ✅ Yes | Last provider that supplied metadata (see below) |
| `last_checked_at` | float | ❌ No (default: `0`) | ✅ Yes | Unix timestamp of last metadata check |
| `hash_status` | string | ❌ No (default: `"completed"`) | ✅ Yes | Hash calculation status: `"pending"`, `"calculating"`, `"completed"`, `"failed"` |
| `autov3` | string\|null | ❌ No | ✅ Yes | CivitAI AutoV3 hash (first 12 chars, lowercase hex) sourced from the safetensors embedded metadata (`sshs_model_hash` / `modelspec.hash_sha256`). **Absent** = not yet checked (may be backfilled later); **`null`** = checked but unavailable (header has no recognized hash); **12-char hex string** = value |
---
@@ -287,6 +288,7 @@ These fields are automatically synchronized with the filesystem:
- `preview_url` — Updated if preview file is moved/removed
- `sha256` — Updated during hash calculation (when `hash_status="pending"`)
- `hash_status` — Updated during hash calculation
- `autov3` — Set when metadata is first created (from safetensors header); may be backfilled later for entries where it is absent
- `last_checked_at` — Timestamp of scan
- `metadata_source` — Set based on metadata provider
@@ -345,6 +347,7 @@ These fields can be edited by users at any time through the Lora Manager UI or b
| `metadata_source` | `null` |
| `last_checked_at` | `0` |
| `hash_status` | `"completed"` |
| `autov3` | absent (not checked) or `null` (checked, no value) |
| `usage_tips` | `"{}"` (LoRA only) |
| `model_type` | `"checkpoint"` or `"embedding"` (not present in LoRA models) |
@@ -354,6 +357,7 @@ These fields can be edited by users at any time through the Lora Manager UI or b
| Version | Date | Changes |
|---------|------|---------|
| 1.1 | 2026-08 | Added `autov3` field (CivitAI AutoV3 hash with three-state semantics) |
| 1.0 | 2026-03 | Initial schema documentation |
---
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+45 -2
View File
@@ -449,6 +449,12 @@
"compact": "7 (1080p), 8 (2K), 10 (4K)"
},
"displayDensityWarning": "Warning: Higher densities may cause performance issues on systems with limited resources.",
"recipesLayout": "Recipes Layout",
"recipesLayoutHelp": "Choose how recipe cards are arranged: a uniform grid or a masonry (Pinterest-style) layout that preserves each image's aspect ratio.",
"recipesLayoutOptions": {
"grid": "Grid",
"masonry": "Masonry"
},
"showFolderSidebar": "Show Folder Sidebar",
"showFolderSidebarHelp": "Toggle the folder navigation sidebar on model pages. When disabled, the sidebar and hover area stay hidden.",
"cardInfoDisplay": "Card Info Display",
@@ -678,6 +684,7 @@
"deepseek": "DeepSeek",
"groq": "Groq",
"openrouter": "OpenRouter",
"google": "Gemini",
"opencode-go": "OpenCode Go",
"custom": "Custom (OpenAI-compatible)"
},
@@ -714,7 +721,9 @@
"versionsCount": "Local Versions",
"versionsCountDesc": "Most versions first",
"versionsCountAsc": "Fewest versions first",
"versionIdDesc": "Newest version first"
"versionIdDesc": "Newest version first",
"random": "Random",
"randomAction": "Randomize (shuffle)"
},
"refresh": {
"title": "Refresh model list",
@@ -771,6 +780,8 @@
"deleteAll": "Delete Selected",
"downloadMissingLoras": "Download Missing LoRAs",
"downloadExamples": "Download Example Images",
"downloadMissingExamples": "Download Missing",
"reprocessExamples": "Re-process All",
"clear": "Clear Selection",
"skipMetadataRefreshCount": "Skip ({count} models)",
"resumeMetadataRefreshCount": "Resume ({count} models)",
@@ -806,6 +817,8 @@
"sendToWorkflowReplace": "Send to Workflow (Replace)",
"openExamples": "Open Examples Folder",
"downloadExamples": "Download Example Images",
"downloadMissingExamples": "Download Missing",
"reprocessExamples": "Re-process All",
"replacePreview": "Replace Preview",
"setContentRating": "Set Content Rating",
"moveToFolder": "Move to Folder",
@@ -1548,6 +1561,7 @@
"empty": "No version history available for this model yet.",
"error": "Failed to load versions.",
"missingModelId": "This model is missing a Civitai model id.",
"hfGroupInfo": "This is a HuggingFace model group. Open the library to see all versions in the grid.",
"confirm": {
"delete": "Delete this version from your library?"
},
@@ -1574,6 +1588,21 @@
"downloadCsv": "Download CSV",
"columnModelName": "Model Name",
"columnError": "Error"
},
"downloadBatchSummary": {
"title": "Batch Download Summary",
"statSuccess": "Success",
"statFailed": "Failed",
"statTotal": "Total",
"successMessage": "All {count} models downloaded successfully",
"completedWithErrors": "Completed with errors",
"failed": "Download failed",
"failedItems": "Failed Items ({count})",
"columnName": "Model Name",
"columnError": "Error",
"close": "Close",
"copyReport": "Copy Report",
"retryFailed": "Retry Failed ({count})"
}
},
"modelTags": {
@@ -1751,6 +1780,12 @@
"checkingMessage": "Please wait while we check for the latest version.",
"showNotifications": "Show update notifications",
"latestBadge": "Latest",
"latestMain": "Latest main",
"channel": "Update Channel",
"channels": {
"release": "Release",
"nightly": "Nightly"
},
"updateProgress": {
"preparing": "Preparing update...",
"installing": "Installing update...",
@@ -1771,6 +1806,15 @@
"warning": "Warning: Nightly builds may contain experimental features and could be unstable.",
"enable": "Enable Nightly Updates"
},
"channelSwitch": {
"nightlyTitle": "Switch to Nightly Channel",
"nightlyMessage": "Switching to Nightly will initialize a Git repository and track the latest main branch commits. Updates will be more frequent but may be unstable. You can switch back to Release at any time.",
"releaseTitle": "Switch to Release Channel",
"releaseMessage": "Switching to Release will checkout the latest stable release tag. You can switch back to Nightly at any time.",
"switching": "Switching to {channel} channel...",
"completed": "Successfully switched to {channel} channel",
"failed": "Failed to switch channel"
},
"banners": {
"recent": "Recent messages",
"empty": "No recent banners yet.",
@@ -2015,7 +2059,6 @@
"presetNameTooLong": "Preset name must be {max} characters or less",
"presetNameInvalidChars": "Preset name contains invalid characters",
"presetNameExists": "A preset with this name already exists",
"maxPresetsReached": "Maximum {max} presets allowed. Delete one to add more.",
"presetNotFound": "Preset not found",
"invalidPreset": "Invalid preset data",
"deletePresetFailed": "Failed to delete preset",
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+2248 -2205
View File
File diff suppressed because it is too large Load Diff
+15 -10
View File
@@ -1,9 +1,13 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import os
import platform
import posixpath
import threading
from pathlib import Path
import folder_paths # type: ignore
import folder_paths # pyright: ignore[reportMissingImports]
from typing import Any, Dict, Iterable, List, Mapping, Optional, Set, Tuple
import logging
import json
@@ -90,7 +94,7 @@ def _resolve_valid_default_root(
def _normalize_folder_paths_for_comparison(
folder_paths: Mapping[str, Iterable[str]],
folder_paths: Mapping[str, Any],
) -> Dict[str, Set[str]]:
"""Normalize folder paths for comparison across libraries."""
@@ -482,7 +486,7 @@ class Config:
import ctypes
FILE_ATTRIBUTE_REPARSE_POINT = 0x400
attrs = ctypes.windll.kernel32.GetFileAttributesW(str(path)) # type: ignore[attr-defined]
attrs = ctypes.windll.kernel32.GetFileAttributesW(str(path)) # pyright: ignore[reportAttributeAccessIssue]
return attrs != -1 and (attrs & FILE_ATTRIBUTE_REPARSE_POINT)
except Exception as e:
logger.error(f"Error checking Windows reparse point: {e}")
@@ -491,7 +495,7 @@ class Config:
logger.error(f"Error checking link status for {path}: {e}")
return False
def _entry_is_symlink(self, entry: os.DirEntry) -> bool:
def _entry_is_symlink(self, entry: os.DirEntry[str]) -> bool:
"""Check if a directory entry is a symlink, including Windows junctions."""
if entry.is_symlink():
return True
@@ -500,7 +504,7 @@ class Config:
import ctypes
FILE_ATTRIBUTE_REPARSE_POINT = 0x400
attrs = ctypes.windll.kernel32.GetFileAttributesW(entry.path) # type: ignore[attr-defined]
attrs = ctypes.windll.kernel32.GetFileAttributesW(entry.path) # pyright: ignore[reportAttributeAccessIssue]
return attrs != -1 and (attrs & FILE_ATTRIBUTE_REPARSE_POINT)
except Exception:
pass
@@ -1126,8 +1130,8 @@ class Config:
def _apply_library_paths(
self,
folder_paths: Mapping[str, Iterable[str]],
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
folder_paths: Mapping[str, Any],
extra_folder_paths: Optional[Mapping[str, Any]] = None,
recipes_path: str = "",
) -> None:
self._path_mappings.clear()
@@ -1432,12 +1436,13 @@ class Config:
# ('_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"
config: Config
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]
config = _sys.modules[_CONFIG_SENTINEL].config
else:
config: Config = Config()
config = Config()
# Register the sentinel so re-imports of py.config find us.
_sentinel_mod = _types.ModuleType(_CONFIG_SENTINEL)
_sentinel_mod.config = config
setattr(_sentinel_mod, "config", config)
_sys.modules[_CONFIG_SENTINEL] = _sentinel_mod
+1 -1
View File
@@ -14,7 +14,7 @@ standalone_mode = (
if not standalone_mode:
setup_logging()
from server import PromptServer # type: ignore
from server import PromptServer # pyright: ignore[reportMissingImports]
from .config import config
from .services.model_service_factory import (
+2 -2
View File
@@ -22,7 +22,7 @@ if not standalone_mode:
logger.info("ComfyUI Metadata Collector initialized")
def get_metadata(prompt_id=None): # type: ignore[no-redef]
def get_metadata(prompt_id=None): # pyright: ignore[reportRedeclaration]
"""Helper function to get metadata from the registry"""
registry = MetadataRegistry()
return registry.get_metadata(prompt_id)
@@ -31,6 +31,6 @@ else:
def init():
logger.info("ComfyUI Metadata Collector disabled in standalone mode")
def get_metadata(prompt_id=None): # type: ignore[no-redef]
def get_metadata(prompt_id=None): # pyright: ignore[reportRedeclaration]
"""Dummy implementation for standalone mode"""
return {}
+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]
+17 -7
View File
@@ -16,7 +16,7 @@ class MetadataHook:
execution = None
try:
# Try direct import first
import execution # type: ignore
import execution # pyright: ignore[reportMissingImports]
except ImportError:
# Try to locate from system modules
for module_name in sys.modules:
@@ -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
+47 -14
View File
@@ -1,7 +1,8 @@
import time
from nodes import NODE_CLASS_MAPPINGS # type: ignore
from typing import Any
from nodes import NODE_CLASS_MAPPINGS # pyright: ignore[reportMissingImports, reportAttributeAccessIssue]
from .node_extractors import NODE_EXTRACTORS, GenericNodeExtractor
from .constants import METADATA_CATEGORIES, IMAGES
from .constants import METADATA_CATEGORIES, IMAGES, OVERWRITE
class MetadataRegistry:
@@ -9,6 +10,15 @@ class MetadataRegistry:
_instance = None
current_prompt_id: Any = None
current_prompt: Any = None
metadata: dict[str, Any] = {}
prompt_metadata: dict[str, Any] = {}
executed_nodes: set[str] = set()
node_cache: dict[str, Any] = {}
max_prompt_history: int = 3
metadata_categories: list[str] = METADATA_CATEGORIES
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
@@ -61,6 +71,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 +86,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 +138,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 +182,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 +204,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)
+97 -5
View File
@@ -2,7 +2,8 @@ import json
import os
import re
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER, OVERWRITE
from .overwrite_utils import collect_overwrite_params
def _store_checkpoint_metadata(metadata, node_id, model_name):
@@ -31,11 +32,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):
@@ -1154,6 +1222,28 @@ 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 = collect_overwrite_params(inputs)
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 = {
@@ -1221,5 +1311,7 @@ NODE_EXTRACTORS = {
"CFGGuider": CFGGuiderExtractor, # Add CFGGuider
# Image
"VAEDecode": VAEDecodeExtractor, # Added VAEDecode extractor
# Metadata overwrite
"MetadataOverwriteLM": MetadataOverwriteExtractor,
# Add other nodes as needed
}
+42
View File
@@ -0,0 +1,42 @@
"""Shared helpers for Metadata Overwrite node metadata collection.
Used by both the MetadataOverwriteLM node (execution time) and the
MetadataOverwriteExtractor (hook time) so the conversion/filtering logic
cannot drift between the two paths.
"""
import logging
from typing import Any, Dict
from ..utils.utils import model_patcher_to_name
from .constants import CLIP_SKIP_SENTINEL, METADATA_OVERWRITE_FIELDS
logger = logging.getLogger(__name__)
def collect_overwrite_params(values: Dict[str, Any]) -> Dict[str, Any]:
"""Convert node input values into non-default overwrite parameters.
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. The ``model`` field accepts either a manual string or
a wired MODEL (ModelPatcher) connection; in the latter case the source
model name is extracted from the patcher's ``cached_patcher_init`` and
stored as a ComfyUI-style relative path.
"""
result: Dict[str, Any] = {}
for key in METADATA_OVERWRITE_FIELDS:
value = values.get(key)
if key == "model" and not isinstance(value, str):
value = model_patcher_to_name(value)
if value is None:
logger.warning(
"Could not extract model name from wired MODEL input "
"(no cached_patcher_init); model metadata overwrite skipped"
)
if key == "clip_skip":
if value != CLIP_SKIP_SENTINEL:
result[key] = value
elif value:
result[key] = value
return result
+2 -2
View File
@@ -43,7 +43,7 @@ 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]:
) -> tuple[Any, 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.
@@ -73,7 +73,7 @@ async def _find_model_entry(
async def _find_scanner_for_model(
model_path: str,
) -> tuple[object, object] | tuple[None, None]:
) -> tuple[Any, object] | tuple[None, None]:
"""Find the (scanner, cache_entry) responsible for *model_path*."""
scanner, entry, _ = await _find_model_entry(model_path)
return scanner, entry
+6 -6
View File
@@ -1,7 +1,7 @@
import logging
from typing import List, Tuple
import comfy.sd # type: ignore
import folder_paths # type: ignore
from typing import Any, List, Tuple
import comfy.sd # pyright: ignore[reportMissingImports]
import folder_paths # pyright: ignore[reportMissingImports]
from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_comfyui
logger = logging.getLogger(__name__)
@@ -18,9 +18,9 @@ class CheckpointLoaderLM:
CATEGORY = "Lora Manager/loaders"
@classmethod
def INPUT_TYPES(s):
def INPUT_TYPES(cls):
# Get list of checkpoint names from scanner (includes extra folder paths)
checkpoint_names = s._get_checkpoint_names()
checkpoint_names = cls._get_checkpoint_names()
return {
"required": {
"ckpt_name": (
@@ -89,7 +89,7 @@ class CheckpointLoaderLM:
logger.error(f"Error getting checkpoint names: {e}")
return []
def load_checkpoint(self, ckpt_name: str) -> Tuple:
def load_checkpoint(self, ckpt_name: str) -> Tuple[Any, Any, Any]:
"""Load a checkpoint by name, supporting extra folder paths
Args:
+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 # pyright: ignore[reportMissingImports] # noqa: C0415
import comfy.utils # pyright: ignore[reportMissingImports] # 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)
+2 -2
View File
@@ -1,8 +1,8 @@
import importlib
import logging
import comfy.sd # type: ignore
import comfy.utils # type: ignore
import comfy.sd # pyright: ignore[reportMissingImports]
import comfy.utils # pyright: ignore[reportMissingImports]
from ..utils.utils import get_lora_info_absolute
from .utils import (
+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) # pyright: ignore[reportAssignmentType]
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,)
+169
View File
@@ -0,0 +1,169 @@
"""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
from ..metadata_collector.overwrite_utils import collect_overwrite_params
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,MODEL",
{
"default": "",
"widgetType": "STRING",
"tooltip": (
"The checkpoint or diffusion model (UNet) used "
"for generation. Fill in the name manually or "
"connect a MODEL output — the model name is then "
"extracted automatically. 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.
The ``model`` field accepts either a manual string or a wired MODEL
(ModelPatcher) connection; in the latter case the underlying model
name is extracted from the patcher's ``cached_patcher_init`` and
stored as a ComfyUI-style relative path.
"""
return (collect_overwrite_params(kwargs),)
+12 -13
View File
@@ -15,15 +15,15 @@ import os
import re
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
from typing import Any, Dict, List, Optional, Tuple, Union, cast
import comfy.utils # type: ignore
import folder_paths # type: ignore
import comfy.utils # pyright: ignore[reportMissingImports]
import folder_paths # pyright: ignore[reportMissingImports]
import torch
import torch.nn as nn
from safetensors import safe_open
from nunchaku.lora.flux.nunchaku_converter import (
from nunchaku.lora.flux.nunchaku_converter import ( # pyright: ignore[reportMissingTypeStubs]
pack_lowrank_weight,
unpack_lowrank_weight,
)
@@ -87,10 +87,6 @@ def _rename_layer_underscore_layer_name(old_name: str) -> str:
return new_name
def _is_indexable_module(module):
return isinstance(module, (nn.ModuleList, nn.Sequential, list, tuple))
def _get_module_by_name(model: nn.Module, name: str) -> Optional[nn.Module]:
if not name:
return model
@@ -100,7 +96,7 @@ def _get_module_by_name(model: nn.Module, name: str) -> Optional[nn.Module]:
continue
if hasattr(module, part):
module = getattr(module, part)
elif part.isdigit() and _is_indexable_module(module):
elif part.isdigit() and isinstance(module, (nn.ModuleList, nn.Sequential, list, tuple)):
try:
module = module[int(part)]
except (IndexError, TypeError):
@@ -267,7 +263,9 @@ def _handle_proj_out_split(lora_dict: Dict[str, Dict[str, torch.Tensor]], base_k
return result, consumed
def _apply_lora_to_module(module: nn.Module, a_tensor: torch.Tensor, b_tensor: torch.Tensor, module_name: str, model: nn.Module) -> None:
def _apply_lora_to_module(module: Any, a_tensor: torch.Tensor, b_tensor: torch.Tensor, module_name: str, model: Any) -> None:
# These modules are dynamic torch containers; monkey-patched attributes
# below are set at runtime, so the module/model types are deliberately Any.
if not hasattr(module, "in_features") or not hasattr(module, "out_features"):
raise ValueError(f"{module_name}: unsupported module without in/out features")
if a_tensor.shape[1] != module.in_features or b_tensor.shape[0] != module.out_features:
@@ -336,7 +334,7 @@ def _apply_lora_to_module(module: nn.Module, a_tensor: torch.Tensor, b_tensor: t
raise ValueError(f"{module_name}: unsupported module type {type(module)}")
def reset_lora_v2(model: nn.Module) -> None:
def reset_lora_v2(model: Any) -> None:
slots = getattr(model, "_lora_slots", None)
if not slots:
return
@@ -344,6 +342,7 @@ def reset_lora_v2(model: nn.Module) -> None:
module = _get_module_by_name(model, name)
if module is None:
continue
module = cast(Any, module)
module_type = info.get("type", "nunchaku")
if module_type == "nunchaku":
base_rank = info["base_rank"]
@@ -371,7 +370,7 @@ def reset_lora_v2(model: nn.Module) -> None:
def compose_loras_v2(model: nn.Module, lora_configs: List[Tuple[Union[str, Path, Dict[str, torch.Tensor]], float]], apply_awq_mod: bool = True) -> bool:
del apply_awq_mod # retained for interface compatibility
reset_lora_v2(model)
aggregated_weights: Dict[str, List[Dict[str, object]]] = defaultdict(list)
aggregated_weights: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
saw_supported_format = False
unresolved_targets = 0
@@ -471,7 +470,7 @@ def compose_loras_v2(model: nn.Module, lora_configs: List[Tuple[Union[str, Path,
class ComfyQwenImageWrapperLM(nn.Module):
def __init__(self, model: nn.Module, config=None, apply_awq_mod: bool = True):
super().__init__()
self.model = model
self.model: Any = model
self.config = {} if config is None else config
self.dtype = next(model.parameters()).dtype
self.loras: List[Tuple[Union[str, Path, Dict[str, torch.Tensor]], float]] = []
+2 -2
View File
@@ -67,7 +67,7 @@ class PromptLM:
stack = inspect.stack()
if len(stack) > 2 and stack[2].function == "get_input_info":
optional_inputs = _PromptOptionalInputs(optional_inputs) # type: ignore[assignment]
optional_inputs = _PromptOptionalInputs(optional_inputs) # pyright: ignore[reportAssignmentType]
return {
"required": {
@@ -126,7 +126,7 @@ class PromptLM:
else:
prompt = expanded_text
from nodes import CLIPTextEncode # type: ignore
from nodes import CLIPTextEncode # pyright: ignore[reportMissingImports, reportAttributeAccessIssue]
conditioning = CLIPTextEncode().encode(clip, prompt)[0]
return (conditioning, prompt)
+362 -130
View File
@@ -5,7 +5,7 @@ import time
import uuid
from typing import Any, Dict, Optional
import numpy as np
import folder_paths # type: ignore
import folder_paths # pyright: ignore[reportMissingImports]
from ..services.service_registry import ServiceRegistry
from ..metadata_collector.metadata_processor import MetadataProcessor
from ..metadata_collector import get_metadata
@@ -13,9 +13,159 @@ from ..utils.constants import CARD_PREVIEW_WIDTH
from ..utils.exif_utils import ExifUtils
from ..utils.utils import calculate_recipe_fingerprint, sanitize_folder_name
from PIL import Image, PngImagePlugin
import piexif
import piexif # pyright: ignore[reportMissingTypeStubs]
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": (
@@ -84,6 +252,13 @@ class SaveImageLM:
"tooltip": "When enabled, embeds generation parameters into the saved image metadata. Disable to skip writing generation metadata.",
},
),
"add_loras_to_prompt": (
"BOOLEAN",
{
"default": False,
"tooltip": "When enabled, appends the LoRA syntax line (e.g. <lora:name:strength>) after the positive prompt in the saved metadata.",
},
),
"add_counter_to_filename": (
"BOOLEAN",
{
@@ -142,148 +317,197 @@ 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, Any], add_loras_to_prompt: bool = False) -> 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[str, Any]] = []
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[str, Any]] = []
if ckpt_civitai.get("id", 0) > 0:
ckpt_resource: dict[str, Any] = {}
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[str, Any] = {"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
prompt_line = prompt if prompt else ""
if add_loras_to_prompt and loras_text:
prompt_line = f"{prompt_line}\n{loras_text}" if prompt_line else loras_text
lines = [prompt_line] if prompt_line 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
@@ -573,10 +797,13 @@ 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,
save_as_recipe=False,
add_loras_to_prompt=False,
):
"""Save images with metadata"""
results = []
@@ -585,7 +812,7 @@ class SaveImageLM:
raw_metadata = get_metadata()
metadata_dict = MetadataProcessor.to_dict(raw_metadata, id)
metadata = self.format_metadata(metadata_dict)
metadata = self.format_metadata(metadata_dict, add_loras_to_prompt)
# Process filename_prefix with pattern substitution
filename_prefix = self.format_filename(filename_prefix, metadata_dict)
@@ -627,15 +854,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}")
@@ -722,10 +948,13 @@ 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,
save_as_recipe=False,
add_loras_to_prompt=False,
):
"""Process and save image with metadata"""
# Make sure the output directory exists
@@ -751,10 +980,13 @@ class SaveImageLM:
extra_pnginfo,
lossless_webp,
quality,
webp_method,
jpeg_subsampling,
embed_workflow,
save_with_metadata,
add_counter_to_filename,
save_as_recipe,
add_loras_to_prompt,
)
return {
+27 -6
View File
@@ -1,12 +1,27 @@
import logging
import os
from typing import List, Tuple
import comfy.sd # type: ignore
from typing import Any, List, Tuple
import comfy.sd # pyright: ignore[reportMissingImports]
from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_comfyui
logger = logging.getLogger(__name__)
def _reload_gguf_unet(
unet_path: str, weight_dtype: str, disable_dynamic: bool = False
) -> object:
"""Reload a GGUF diffusion model from disk (cached_patcher_init factory).
Mirrors the GGUF branch of UNETLoaderLM.load_unet so ModelPatcher
deepclone/dynamic machinery can rebuild GGUF models with the correct
GGMLOps. ``disable_dynamic`` is accepted for signature compatibility
with core ComfyUI loaders.
"""
loader = UNETLoaderLM()
model, = loader._load_gguf_unet(unet_path, unet_path, weight_dtype)
return model
class UNETLoaderLM:
"""UNET Loader with support for extra folder paths
@@ -19,9 +34,9 @@ class UNETLoaderLM:
CATEGORY = "Lora Manager/loaders"
@classmethod
def INPUT_TYPES(s):
def INPUT_TYPES(cls):
# Get list of unet names from scanner (includes extra folder paths)
unet_names = s._get_unet_names()
unet_names = cls._get_unet_names()
return {
"required": {
"unet_name": (
@@ -90,7 +105,7 @@ class UNETLoaderLM:
logger.error(f"Error getting unet names: {e}")
return []
def load_unet(self, unet_name: str, weight_dtype: str) -> Tuple:
def load_unet(self, unet_name: str, weight_dtype: str) -> Tuple[Any, ...]:
"""Load a diffusion model by name, supporting extra folder paths
Args:
@@ -133,7 +148,7 @@ class UNETLoaderLM:
def _load_gguf_unet(
self, unet_path: str, unet_name: str, weight_dtype: str
) -> Tuple:
) -> Tuple[Any, ...]:
"""Load a GGUF format diffusion model
Args:
@@ -196,6 +211,12 @@ class UNETLoaderLM:
# Wrap with GGUFModelPatcher
model = GGUFModelPatcher.clone(model)
# Register a reload factory so the MODEL carries its source path
# (cached_patcher_init) like core ComfyUI loaders do — required
# for model-name extraction downstream and for ModelPatcher
# deepclone/dynamic machinery.
model.cached_patcher_init = (_reload_gguf_unet, (unet_path, weight_dtype))
return (model,)
except Exception as e:
+7 -3
View File
@@ -1,3 +1,6 @@
from typing import Any
class AnyType(str):
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
@@ -6,7 +9,7 @@ class AnyType(str):
# Credit to Regis Gaughan, III (rgthree)
class FlexibleOptionalInputType(dict):
class FlexibleOptionalInputType(dict[str, Any]):
"""A special class to make flexible nodes that pass data to our python handlers.
Enables both flexible/dynamic input types (like for Any Switch) or a dynamic number of inputs
@@ -23,6 +26,7 @@ class FlexibleOptionalInputType(dict):
"""
def __init__(self, type):
super().__init__()
self.type = type
def __getitem__(self, key):
@@ -40,7 +44,7 @@ import re
import logging
import copy
import sys
import folder_paths # type: ignore
import folder_paths # pyright: ignore[reportMissingImports]
logger = logging.getLogger(__name__)
@@ -70,7 +74,7 @@ def extract_lora_name(lora_path):
return apply_lora_syntax_format(name_no_ext)
def parse_lora_syntax(text: str) -> list[dict]:
def parse_lora_syntax(text: str) -> list[dict[str, Any]]:
"""Parse <lora:name:strength> syntax from text input into a list of dicts.
Each entry contains: name, model_strength, clip_strength.
+17 -5
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""Base classes for recipe parsers."""
import json
@@ -38,7 +42,7 @@ class RecipeMetadataParser(ABC):
pass
@staticmethod
async def populate_lora_from_civitai(lora_entry: Dict[str, Any], civitai_info_tuple: Tuple[Dict[str, Any], Optional[str]],
async def populate_lora_from_civitai(lora_entry: Dict[str, Any], civitai_info_tuple: Tuple[Dict[str, Any] | None, str | None] | Dict[str, Any],
recipe_scanner=None, base_model_counts=None, hash_value=None) -> Optional[Dict[str, Any]]:
"""
Populate a lora entry with information from Civitai API response
@@ -175,10 +179,18 @@ class RecipeMetadataParser(ABC):
lora_entry['localPath'] = local_path
lora_entry['file_name'] = os.path.splitext(os.path.basename(local_path))[0]
# Get thumbnail from local preview if available
# Get thumbnail from local preview if available.
# Match the cache item by local path first (get_path_by_hash
# cascade: 10-char autov2 / 12-char autov3), then by hash.
lora_cache = await lora_scanner.get_cached_data()
lora_item = next((item for item in lora_cache.raw_data
if item['sha256'].lower() == lora_entry['hash'].lower()), None)
h = (lora_entry.get("hash") or "").lower()
lora_item = next((item for item in lora_cache.raw_data
if (item.get("file_path") or "") == local_path), None)
if lora_item is None:
lora_item = next((item for item in lora_cache.raw_data
if (item.get("sha256") or "").lower() == h
or (item.get("autov3") or "").lower() == h
or (item.get("sha256") or "")[:10].lower() == h), None)
if lora_item and 'preview_url' in lora_item:
lora_entry['thumbnailUrl'] = config.get_preview_static_url(lora_item['preview_url'])
except Exception as e:
@@ -194,7 +206,7 @@ class RecipeMetadataParser(ABC):
return lora_entry
@staticmethod
async def populate_checkpoint_from_civitai(checkpoint: Dict[str, Any], civitai_info: Dict[str, Any]) -> Dict[str, Any]:
async def populate_checkpoint_from_civitai(checkpoint: Dict[str, Any], civitai_info: Dict[str, Any] | Tuple[Dict[str, Any] | None, str | None] | None) -> Dict[str, Any]:
"""
Populate checkpoint information from Civitai API response
+4
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import logging
import json
import os
+3 -1
View File
@@ -1,6 +1,7 @@
"""Factory for creating recipe metadata parsers."""
import logging
from typing import Any
from .parsers import (
RecipeFormatParser,
ComfyMetadataParser,
@@ -31,7 +32,8 @@ class RecipeParserFactory:
# First, try CivitaiApiMetadataParser for dict input
if isinstance(metadata, dict):
try:
if CivitaiApiMetadataParser().is_metadata_matching(metadata):
user_comment: Any = metadata
if CivitaiApiMetadataParser().is_metadata_matching(user_comment):
return CivitaiApiMetadataParser()
except Exception as e:
logger.debug(f"CivitaiApiMetadataParser check failed: {e}")
+1 -1
View File
@@ -52,7 +52,7 @@ class AutomaticMetadataParser(RecipeMetadataParser):
negative_and_params = ""
# Initialize metadata
metadata = {
metadata: Dict[str, Any] = {
"prompt": prompt,
"loras": []
}
+117 -56
View File
@@ -4,7 +4,7 @@ import json
import logging
from typing import Dict, Any, Union
from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS
from ..constants import GEN_PARAM_KEYS, VALID_LORA_TYPES
from ...services.metadata_service import get_default_metadata_provider
from ...config import config
@@ -14,15 +14,16 @@ logger = logging.getLogger(__name__)
class CivitaiApiMetadataParser(RecipeMetadataParser):
"""Parser for Civitai image metadata format"""
def is_metadata_matching(self, metadata) -> bool:
def is_metadata_matching(self, user_comment) -> bool:
"""Check if the metadata matches the Civitai image metadata format
Args:
metadata: The metadata from the image (dict)
user_comment: The metadata from the image (dict)
Returns:
bool: True if this parser can handle the metadata
"""
metadata = user_comment
if not metadata or not isinstance(metadata, dict):
return False
@@ -73,7 +74,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
return False
async def parse_metadata( # type: ignore[override]
async def parse_metadata( # pyright: ignore[reportIncompatibleMethodOverride]
self, user_comment, recipe_scanner=None, civitai_client=None,
local_cache: dict[str, Any] | None = None,
) -> Dict[str, Any]:
@@ -89,8 +90,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
Returns:
Dict containing parsed recipe data
"""
metadata: Dict[str, Any] = user_comment # type: ignore[assignment]
metadata = user_comment
metadata: Dict[str, Any] = user_comment
try:
# Get metadata provider instead of using civitai_client directly
metadata_provider = await get_default_metadata_provider()
@@ -116,7 +116,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
metadata = inner_meta
# Initialize result structure
result = {
result: Dict[str, Any] = {
"base_model": None,
"loras": [],
"model": None,
@@ -125,10 +125,10 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
}
# Track already added LoRAs to prevent duplicates
added_loras = {} # key: model_version_id or hash, value: index in result["loras"]
added_loras: Dict[str, Any] = {} # key: model_version_id or hash, value: index in result["loras"]
# Extract hash information from hashes field for LoRA matching
lora_hashes = {}
lora_hashes: Dict[str, Any] = {}
if "hashes" in metadata and isinstance(metadata["hashes"], dict):
for key, hash_value in metadata["hashes"].items():
key_str = str(key)
@@ -184,7 +184,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
if model_info:
result["base_model"] = model_info.get("baseModel", "")
base_model_counts = {}
base_model_counts: Dict[str, int] = {}
# Process standard resources array
if "resources" in metadata and isinstance(metadata["resources"], list):
@@ -196,7 +196,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
# identification because it has an explicit type field and hash,
# unlike modelVersionIds which is a flat list with no type info.
if resource_type == "model":
checkpoint_entry = {
checkpoint_entry: Dict[str, Any] = {
"id": 0,
"modelId": 0,
"name": resource.get("name", "Unknown Model"),
@@ -216,7 +216,8 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
# Try to look up base model from the checkpoint hash
cp_hash = checkpoint_entry.get("hash")
if cp_hash and metadata_provider:
local_cached = local_cache.get(cp_hash) if local_cache else None
# local_cache keys are stored lowercase
local_cached = local_cache.get(cp_hash.lower()) if local_cache else None
if local_cached:
self._populate_entry_from_cache(
checkpoint_entry, local_cached
@@ -294,8 +295,15 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
# Try to get info from Civitai if hash is available
if lora_hash and metadata_provider:
local_cached = local_cache.get(lora_hash) if local_cache else None
# local_cache keys are stored lowercase
local_cached = local_cache.get(lora_hash.lower()) if local_cache else None
if local_cached:
cached_type = self._cache_item_model_type(local_cached)
if cached_type and cached_type not in VALID_LORA_TYPES:
logger.debug(
f"Skipping non-LoRA cache item for hash {lora_hash}"
)
continue
self._populate_entry_from_cache(
lora_entry, local_cached
)
@@ -304,6 +312,12 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
added_loras[str(lora_entry["id"])] = len(
result["loras"]
)
# Mirror base.py:150-151 counts for API-path loras
bm = local_cached.get("base_model") or ""
if bm:
base_model_counts[bm] = base_model_counts.get(
bm, 0
) + 1
else:
try:
civitai_info = (
@@ -649,30 +663,47 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
}
if metadata_provider:
try:
civitai_info = await metadata_provider.get_model_by_hash(
lora_hash
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
lora_hash,
)
if populated_entry is None:
# local_cache keys are stored lowercase
local_cached = local_cache.get(lora_hash.lower()) if local_cache else None
if local_cached:
cached_type = self._cache_item_model_type(local_cached)
if cached_type and cached_type not in VALID_LORA_TYPES:
logger.debug(
f"Skipping non-LoRA cache item for hash {lora_hash}"
)
continue
lora_entry = populated_entry
self._populate_entry_from_cache(lora_entry, local_cached)
# Mirror base.py:150-151 counts for API-path loras
bm = local_cached.get("base_model") or ""
if bm:
base_model_counts[bm] = base_model_counts.get(bm, 0) + 1
if "id" in lora_entry and lora_entry["id"]:
added_loras[str(lora_entry["id"])] = len(result["loras"])
except Exception as e:
logger.error(
f"Error fetching Civitai info for LoRA hash {lora_hash}: {e}"
)
else:
try:
civitai_info = await metadata_provider.get_model_by_hash(
lora_hash
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
lora_hash,
)
if populated_entry is None:
continue
lora_entry = populated_entry
if "id" in lora_entry and lora_entry["id"]:
added_loras[str(lora_entry["id"])] = len(result["loras"])
except Exception as e:
logger.error(
f"Error fetching Civitai info for LoRA hash {lora_hash}: {e}"
)
added_loras[lora_hash] = len(result["loras"])
result["loras"].append(lora_entry)
@@ -711,32 +742,51 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
# Try to get info from Civitai if hash is available
if lora_entry["hash"] and metadata_provider:
try:
civitai_info = await metadata_provider.get_model_by_hash(
lora_hash
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
lora_hash,
)
if populated_entry is None:
# local_cache keys are stored lowercase
local_cached = local_cache.get(lora_hash.lower()) if local_cache else None
if local_cached:
cached_type = self._cache_item_model_type(local_cached)
if cached_type and cached_type not in VALID_LORA_TYPES:
logger.debug(
f"Skipping non-LoRA cache item for hash {lora_hash}"
)
lora_index += 1
continue # Skip invalid LoRA types
lora_entry = populated_entry
continue # Skip non-LoRA cache items
self._populate_entry_from_cache(lora_entry, local_cached)
# Mirror base.py:150-151 counts for API-path loras
bm = local_cached.get("base_model") or ""
if bm:
base_model_counts[bm] = base_model_counts.get(bm, 0) + 1
# If we have a version ID from Civitai, track it for deduplication
if "id" in lora_entry and lora_entry["id"]:
added_loras[str(lora_entry["id"])] = len(result["loras"])
except Exception as e:
logger.error(
f"Error fetching Civitai info for LoRA hash {lora_entry['hash']}: {e}"
)
else:
try:
civitai_info = await metadata_provider.get_model_by_hash(
lora_hash
)
populated_entry = await self.populate_lora_from_civitai(
lora_entry,
civitai_info,
recipe_scanner,
base_model_counts,
lora_hash,
)
if populated_entry is None:
lora_index += 1
continue # Skip invalid LoRA types
lora_entry = populated_entry
# If we have a version ID from Civitai, track it for deduplication
if "id" in lora_entry and lora_entry["id"]:
added_loras[str(lora_entry["id"])] = len(result["loras"])
except Exception as e:
logger.error(
f"Error fetching Civitai info for LoRA hash {lora_entry['hash']}: {e}"
)
# Track by hash if we have it
if lora_hash:
@@ -795,3 +845,14 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
base_model = cache_item.get("base_model", "")
if base_model:
entry["baseModel"] = base_model
@staticmethod
def _cache_item_model_type(cache_item: dict[str, Any]) -> str:
"""Lowercased civitai.model.type of a cache item, or '' when unknown."""
civ = cache_item.get("civitai")
if not isinstance(civ, dict):
return ""
model_info = civ.get("model")
if not isinstance(model_info, dict):
return ""
return (model_info.get("type") or "").lower()
+1 -1
View File
@@ -30,7 +30,7 @@ class MetaFormatParser(RecipeMetadataParser):
prompt = parts[0].strip()
# Initialize metadata
metadata = {"prompt": prompt, "loras": []}
metadata: Dict[str, Any] = {"prompt": prompt, "loras": []}
# Extract negative prompt and parameters if available
if len(parts) > 1:
+10 -2
View File
@@ -91,7 +91,15 @@ class RecipeFormatParser(RecipeMetadataParser):
exists_locally = lora_scanner.has_hash(lora['hash'])
if exists_locally:
lora_cache = await lora_scanner.get_cached_data()
lora_item = next((item for item in lora_cache.raw_data if item['sha256'].lower() == lora['hash'].lower()), None)
# Cascade match: full sha256, stored autov3, or autov2 (sha256[:10]).
h = (lora.get('hash') or '').lower()
lora_item = next(
(item for item in lora_cache.raw_data
if (item.get("sha256") or "").lower() == h
or (item.get("autov3") or "").lower() == h
or (item.get("sha256") or "")[:10].lower() == h),
None
)
if lora_item:
lora_entry['existsLocally'] = True
lora_entry['inLibrary'] = True
@@ -148,7 +156,7 @@ class RecipeFormatParser(RecipeMetadataParser):
checkpoint_data = recipe_metadata.get('checkpoint') or {}
if isinstance(checkpoint_data, dict) and checkpoint_data:
version_id = checkpoint_data.get('modelVersionId') or checkpoint_data.get('id')
checkpoint_entry = {
checkpoint_entry: Dict[str, Any] = {
'id': version_id or 0,
'modelId': checkpoint_data.get('modelId', 0),
'name': checkpoint_data.get('name', 'Unknown Checkpoint'),
+7 -7
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import logging
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Callable, Dict, Mapping
from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Mapping
import jinja2
from aiohttp import web
@@ -84,7 +84,7 @@ class BaseModelRoutes(ABC):
self.metadata_progress_callback = WebSocketBroadcastCallback()
self._handler_set: ModelHandlerSet | None = None
self._handler_mapping: Dict[str, Callable[[web.Request], web.StreamResponse]] | None = None
self._handler_mapping: Dict[str, Callable[[web.Request], Awaitable[web.Response]]] | None = None
self._preview_service = PreviewAssetService(
metadata_manager=MetadataManager,
@@ -131,7 +131,7 @@ class BaseModelRoutes(ABC):
self._handler_set = None
self._handler_mapping = None
def _ensure_handler_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
def _ensure_handler_mapping(self) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
if self._handler_mapping is None:
handler_set = self._create_handler_set()
self._handler_set = handler_set
@@ -220,7 +220,7 @@ class BaseModelRoutes(ABC):
)
@property
def route_handlers(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
def route_handlers(self) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
return self._ensure_handler_mapping()
def setup_routes(self, app: web.Application, prefix: str) -> None:
@@ -237,7 +237,7 @@ class BaseModelRoutes(ABC):
"""Setup model-specific routes."""
raise NotImplementedError
def _parse_specific_params(self, request: web.Request) -> Dict:
def _parse_specific_params(self, request: web.Request) -> Dict[str, Any]:
"""Parse model-specific parameters - to be overridden by subclasses."""
return {}
@@ -253,7 +253,7 @@ class BaseModelRoutes(ABC):
"""Find the appropriate model file from the files list - can be overridden by subclasses."""
return next((file for file in files if file.get("type") in ("Model", "Diffusion Model") and file.get("primary") is True), None)
def get_handler(self, name: str) -> Callable[[web.Request], web.StreamResponse]:
def get_handler(self, name: str) -> Callable[[web.Request], Awaitable[web.StreamResponse]]:
"""Expose handlers for subclasses or tests."""
return self._ensure_handler_mapping()[name]
@@ -285,7 +285,7 @@ class BaseModelRoutes(ABC):
)
return self.model_lifecycle_service
def _make_handler_proxy(self, name: str) -> Callable[[web.Request], web.StreamResponse]:
def _make_handler_proxy(self, name: str) -> Callable[[web.Request], Awaitable[web.StreamResponse]]:
async def proxy(request: web.Request) -> web.StreamResponse:
try:
handler = self.get_handler(name)
+13 -9
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
import logging
import os
from typing import Callable, Mapping
from typing import Awaitable, Callable, Mapping
import jinja2
from aiohttp import web
@@ -61,7 +61,9 @@ class BaseRecipeRoutes:
self._i18n_registered = False
self._startup_hooks_registered = False
self._handler_set: RecipeHandlerSet | None = None
self._handler_mapping: dict[str, Callable] | None = None
self._handler_mapping: Mapping[
str, Callable[[web.Request], Awaitable[web.StreamResponse]]
] | None = None
async def attach_dependencies(self, app: web.Application | None = None) -> None:
"""Resolve shared services from the registry."""
@@ -84,7 +86,9 @@ class BaseRecipeRoutes:
app.on_startup.append(self.attach_dependencies)
self._startup_hooks_registered = True
def to_route_mapping(self) -> Mapping[str, Callable]:
def to_route_mapping(
self,
) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
"""Return a mapping of handler name to coroutine for registrar binding."""
if self._handler_mapping is None:
@@ -124,17 +128,17 @@ class BaseRecipeRoutes:
or os.environ.get("HF_HUB_DISABLE_TELEMETRY", "0") == "0"
)
if not standalone_mode:
from ..metadata_collector import get_metadata # type: ignore[import-not-found]
from ..metadata_collector.metadata_processor import ( # type: ignore[import-not-found]
from ..metadata_collector import get_metadata # pyright: ignore[reportMissingImports]
from ..metadata_collector.metadata_processor import ( # pyright: ignore[reportMissingImports]
MetadataProcessor,
)
from ..metadata_collector.metadata_registry import ( # type: ignore[import-not-found]
from ..metadata_collector.metadata_registry import ( # pyright: ignore[reportMissingImports]
MetadataRegistry,
)
else: # pragma: no cover - optional dependency path
get_metadata = None # type: ignore[assignment]
MetadataProcessor = None # type: ignore[assignment]
MetadataRegistry = None # type: ignore[assignment]
get_metadata = None # pyright: ignore[reportAssignmentType]
MetadataProcessor = None # pyright: ignore[reportAssignmentType]
MetadataRegistry = None # pyright: ignore[reportAssignmentType]
analysis_service = RecipeAnalysisService(
exif_utils=ExifUtils,
+9 -9
View File
@@ -1,5 +1,5 @@
import logging
from typing import Dict, List, Set
from typing import Any, Dict, List, Set
from aiohttp import web
from .base_model_routes import BaseModelRoutes
@@ -28,13 +28,13 @@ class CheckpointRoutes(BaseModelRoutes):
# Attach service dependencies
self.attach_service(self.service)
def setup_routes(self, app: web.Application):
def setup_routes(self, app: web.Application, prefix: str = "checkpoints"):
"""Setup Checkpoint routes"""
# Schedule service initialization on app startup
app.on_startup.append(lambda _: self.initialize_services())
# Setup common routes with 'checkpoints' prefix (includes page route)
super().setup_routes(app, 'checkpoints')
super().setup_routes(app, prefix)
def setup_specific_routes(self, registrar: ModelRouteRegistrar, prefix: str):
"""Setup Checkpoint-specific routes"""
@@ -53,9 +53,9 @@ class CheckpointRoutes(BaseModelRoutes):
"""Get expected model types string for error messages"""
return "Checkpoint"
def _parse_specific_params(self, request: web.Request) -> Dict:
def _parse_specific_params(self, request: web.Request) -> Dict[str, Any]:
"""Parse Checkpoint-specific parameters"""
params: Dict = {}
params: Dict[str, Any] = {}
if 'checkpoint_hash' in request.query:
params['hash_filters'] = {'single_hash': request.query['checkpoint_hash'].lower()}
@@ -70,7 +70,7 @@ class CheckpointRoutes(BaseModelRoutes):
"""Get detailed information for a specific checkpoint by name"""
try:
name = request.match_info.get('name', '')
checkpoint_info = await self.service.get_model_info_by_name(name)
checkpoint_info = await self.service.get_model_info_by_name(name) # pyright: ignore[reportAttributeAccessIssue]
if checkpoint_info:
return web.json_response(checkpoint_info)
@@ -89,7 +89,7 @@ class CheckpointRoutes(BaseModelRoutes):
roots.extend(config.checkpoints_roots or [])
roots.extend(config.extra_checkpoints_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root and root not in seen:
@@ -114,7 +114,7 @@ class CheckpointRoutes(BaseModelRoutes):
roots.extend(config.unet_roots or [])
roots.extend(config.extra_unet_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root and root not in seen:
+4 -4
View File
@@ -26,13 +26,13 @@ class EmbeddingRoutes(BaseModelRoutes):
# Attach service dependencies
self.attach_service(self.service)
def setup_routes(self, app: web.Application):
def setup_routes(self, app: web.Application, prefix: str = "embeddings"):
"""Setup Embedding routes"""
# Schedule service initialization on app startup
app.on_startup.append(lambda _: self.initialize_services())
# Setup common routes with 'embeddings' prefix (includes page route)
super().setup_routes(app, 'embeddings')
super().setup_routes(app, prefix)
def setup_specific_routes(self, registrar: ModelRouteRegistrar, prefix: str):
"""Setup Embedding-specific routes"""
@@ -51,7 +51,7 @@ class EmbeddingRoutes(BaseModelRoutes):
"""Get detailed information for a specific embedding by name"""
try:
name = request.match_info.get('name', '')
embedding_info = await self.service.get_model_info_by_name(name)
embedding_info = await self.service.get_model_info_by_name(name) # pyright: ignore[reportAttributeAccessIssue]
if embedding_info:
return web.json_response(embedding_info)
+8 -4
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import logging
from typing import Callable, Mapping
from typing import Any, Awaitable, Callable, Mapping
from aiohttp import web
@@ -35,7 +35,7 @@ class ExampleImagesRoutes:
*,
ws_manager,
download_manager: DownloadManager | None = None,
processor=ExampleImagesProcessor,
processor: Any = ExampleImagesProcessor,
file_manager=ExampleImagesFileManager,
cleanup_service: ExampleImagesCleanupService | None = None,
) -> None:
@@ -46,7 +46,9 @@ class ExampleImagesRoutes:
self._file_manager = file_manager
self._cleanup_service = cleanup_service or ExampleImagesCleanupService()
self._handler_set: ExampleImagesHandlerSet | None = None
self._handler_mapping: Mapping[str, Callable[[web.Request], web.StreamResponse]] | None = None
self._handler_mapping: Mapping[
str, Callable[[web.Request], Awaitable[web.StreamResponse]]
] | None = None
@classmethod
def setup_routes(cls, app: web.Application, *, ws_manager) -> None:
@@ -61,7 +63,9 @@ class ExampleImagesRoutes:
registrar = ExampleImagesRouteRegistrar(app)
registrar.register_routes(self.to_route_mapping())
def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
def to_route_mapping(
self,
) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
"""Return the registrar-compatible mapping of handler names to callables."""
if self._handler_mapping is None:
@@ -3,7 +3,7 @@ from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Callable, Mapping
from typing import Awaitable, Callable, Mapping
from aiohttp import web
@@ -170,7 +170,7 @@ class ExampleImagesHandlerSet:
management: ExampleImagesManagementHandler
files: ExampleImagesFileHandler
def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
"""Flatten handler methods into the registrar mapping."""
return {
+95 -48
View File
@@ -276,7 +276,7 @@ def _collect_comfyui_session_logs(
) -> dict[str, Any]:
if log_entries is None:
try:
import app.logger as comfy_logger
import app.logger as comfy_logger # pyright: ignore[reportMissingImports]
log_entries = list(comfy_logger.get_logs() or [])
except Exception as exc: # pragma: no cover - environment dependent
@@ -422,10 +422,10 @@ class PromptServerProtocol(Protocol):
"""Subset of PromptServer used by the handlers."""
instance: "PromptServerProtocol"
sockets: dict # maps clientId (sid) → WebSocketResponse
sockets: dict[str, Any] # maps clientId (sid) → WebSocketResponse
def send_sync(
self, event: str, payload: dict | None = None, sid: str | None = None
self, event: str, payload: dict[str, Any] | None = None, sid: str | None = None
) -> None: # pragma: no cover - protocol
...
@@ -443,7 +443,12 @@ class UsageStatsFactory(Protocol):
class MetadataProviderProtocol(Protocol):
async def get_model_versions(
self, model_id: int
) -> dict | None: # pragma: no cover - protocol
) -> dict[str, Any] | None: # pragma: no cover - protocol
...
async def get_user_models(
self, username: str, cursor: str | None = None
) -> Any: # pragma: no cover - protocol
...
@@ -466,16 +471,16 @@ class MetadataArchiveManagerProtocol(Protocol):
class BackupServiceProtocol(Protocol):
async def create_snapshot(
self, *, snapshot_type: str = "manual", persist: bool = False
) -> dict: # pragma: no cover - protocol
) -> dict[str, Any]: # pragma: no cover - protocol
...
async def restore_snapshot(self, archive_path: str) -> dict: # pragma: no cover - protocol
async def restore_snapshot(self, archive_path: str) -> dict[str, Any]: # pragma: no cover - protocol
...
def get_status(self) -> dict: # pragma: no cover - protocol
def get_status(self) -> dict[str, Any]: # pragma: no cover - protocol
...
def get_available_snapshots(self) -> list[dict]: # pragma: no cover - protocol
def get_available_snapshots(self) -> list[dict[str, Any]]: # pragma: no cover - protocol
...
@@ -491,7 +496,7 @@ class NodeRegistry:
def __init__(self) -> None:
self._lock = asyncio.Lock()
# sid → {unique_id → node_info}
self._tab_nodes: Dict[str, Dict[str, dict]] = {}
self._tab_nodes: Dict[str, Dict[str, dict[str, Any]]] = {}
self._ready = asyncio.Event()
self._waiting_clients: set[str] = set()
@@ -504,7 +509,7 @@ class NodeRegistry:
# Helpers to build one node dict (extracted so it's reused for each tab)
# ------------------------------------------------------------------
@staticmethod
def _build_node_dict(node: dict) -> dict:
def _build_node_dict(node: dict[str, Any]) -> dict[str, Any]:
node_id = node["node_id"]
graph_id = str(node["graph_id"])
unique_id = f"{graph_id}:{node_id}"
@@ -513,11 +518,11 @@ class NodeRegistry:
bgcolor = node.get("bgcolor") or DEFAULT_NODE_COLOR
raw_capabilities = node.get("capabilities")
capabilities: dict = {}
capabilities: dict[str, Any] = {}
if isinstance(raw_capabilities, dict):
capabilities = dict(raw_capabilities)
raw_widget_names: list | None = node.get("widget_names")
raw_widget_names: list[Any] | None = node.get("widget_names")
if not isinstance(raw_widget_names, list):
capability_widget_names = capabilities.get("widget_names")
raw_widget_names = (
@@ -565,9 +570,9 @@ class NodeRegistry:
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
async def register_nodes(self, sid: str, nodes: list[dict]) -> None:
async def register_nodes(self, sid: str, nodes: list[dict[str, Any]]) -> None:
"""Register/replace the node list for a single ComfyUI tab (identified by *sid*)."""
tab_nodes: dict[str, dict] = {}
tab_nodes: dict[str, dict[str, Any]] = {}
for node in nodes:
nd = self._build_node_dict(node)
tab_nodes[nd["unique_id"]] = nd
@@ -602,7 +607,7 @@ class NodeRegistry:
except asyncio.TimeoutError:
return False
async def get_merged_registry(self, active_sids: set[str] | None = None) -> dict:
async def get_merged_registry(self, active_sids: set[str] | None = None) -> dict[str, Any]:
"""Return the union of all known tab nodes, pruning any tab that is no
longer connected."""
async with self._lock:
@@ -619,8 +624,8 @@ class NodeRegistry:
len(stale_sids), stale_sids,
)
merged: dict[str, dict] = {}
tab_info: dict[str, dict] = {}
merged: dict[str, dict[str, Any]] = {}
tab_info: dict[str, dict[str, Any]] = {}
for sid, nodes in self._tab_nodes.items():
tab_info[sid] = {
"node_count": len(nodes),
@@ -653,7 +658,7 @@ class SupportersHandler:
def __init__(self, logger: logging.Logger | None = None) -> None:
self._logger = logger or logging.getLogger(__name__)
def _load_supporters(self) -> dict:
def _load_supporters(self) -> dict[str, Any]:
"""Load supporters data from JSON file."""
try:
current_file = os.path.abspath(__file__)
@@ -1229,10 +1234,8 @@ class DoctorHandler:
settings_snapshot = _sanitize_sensitive_data(
getattr(self._settings, "settings", {}) or {}
)
startup_messages_getter = getattr(self._settings, "get_startup_messages", None)
startup_messages = (
list(startup_messages_getter()) if callable(startup_messages_getter) else []
)
startup_messages_getter: Any = getattr(self._settings, "get_startup_messages", None)
startup_messages = list(startup_messages_getter()) if startup_messages_getter else []
environment = {
"app_version": app_version,
@@ -1439,7 +1442,7 @@ class SettingsHandler:
*,
settings_service=None,
metadata_provider_updater: Callable[
[], Awaitable[None]
[], Awaitable[Any]
] = update_metadata_providers,
downloader_factory: Callable[
[], Awaitable[DownloaderProtocol]
@@ -1484,8 +1487,8 @@ class SettingsHandler:
settings_file = getattr(self._settings, "settings_file", None)
if settings_file:
response_data["settings_file"] = settings_file
messages_getter = getattr(self._settings, "get_startup_messages", None)
messages = list(messages_getter()) if callable(messages_getter) else []
messages_getter: Any = getattr(self._settings, "get_startup_messages", None)
messages = list(messages_getter()) if messages_getter else []
return web.json_response(
{
"success": True,
@@ -1562,6 +1565,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",
@@ -2000,11 +2008,11 @@ async def _noop_backup_service() -> None:
@dataclass
class ServiceRegistryAdapter:
get_lora_scanner: Callable[[], Awaitable]
get_checkpoint_scanner: Callable[[], Awaitable]
get_embedding_scanner: Callable[[], Awaitable]
get_downloaded_version_history_service: Callable[[], Awaitable]
get_backup_service: Callable[[], Awaitable] = _noop_backup_service
get_lora_scanner: Callable[[], Awaitable[Any]]
get_checkpoint_scanner: Callable[[], Awaitable[Any]]
get_embedding_scanner: Callable[[], Awaitable[Any]]
get_downloaded_version_history_service: Callable[[], Awaitable[Any]]
get_backup_service: Callable[[], Awaitable[Any]] = _noop_backup_service
class ModelLibraryHandler:
@@ -2045,8 +2053,8 @@ class ModelLibraryHandler:
return await self._service_registry.get_downloaded_version_history_service()
@staticmethod
def _with_downloaded_flag(versions: list[dict]) -> list[dict]:
enriched: list[dict] = []
def _with_downloaded_flag(versions: list[dict[str, Any]]) -> list[dict[str, Any]]:
enriched: list[dict[str, Any]] = []
for version in versions:
entry = dict(version)
entry.setdefault("hasBeenDownloaded", True)
@@ -2239,7 +2247,7 @@ class ModelLibraryHandler:
checkpoint_scanner = await self._service_registry.get_checkpoint_scanner()
embedding_scanner = await self._service_registry.get_embedding_scanner()
results: list[dict] = []
results: list[dict[str, Any]] = []
for model_id in model_ids:
lora_versions = await lora_scanner.get_model_versions_by_id(model_id)
if lora_versions:
@@ -2348,7 +2356,7 @@ class ModelLibraryHandler:
)
try:
model_version_id = int(data.get("modelVersionId"))
model_version_id = int(data.get("modelVersionId")) # pyright: ignore[reportArgumentType]
except (TypeError, ValueError):
return web.json_response(
{"success": False, "error": "Parameter modelVersionId must be an integer"},
@@ -2460,10 +2468,11 @@ class ModelLibraryHandler:
"checkpoint": checkpoint_scanner,
"embedding": embedding_scanner,
}
scanner = scanner_map.get(found_type)
scanner = scanner_map.get(found_type or "")
if scanner:
persist = getattr(scanner, "_persist_current_cache", None)
if callable(persist):
scanner.bump_cache_version()
persist: Any = getattr(scanner, "_persist_current_cache", None)
if persist:
await persist()
history_service = await self._get_download_history_service()
@@ -2585,6 +2594,8 @@ class ModelLibraryHandler:
status=400,
)
cursor = request.query.get("cursor")
metadata_provider = await self._metadata_provider_factory()
if not metadata_provider:
return web.json_response(
@@ -2593,7 +2604,7 @@ class ModelLibraryHandler:
)
try:
models = await metadata_provider.get_user_models(username)
result = await metadata_provider.get_user_models(username, cursor)
except NotImplementedError:
return web.json_response(
{
@@ -2603,14 +2614,35 @@ class ModelLibraryHandler:
status=501,
)
if models is None:
if result is None:
return web.json_response(
{"success": False, "error": "Failed to fetch user models"},
status=502,
)
if isinstance(result, dict):
models = result.get("items")
next_cursor = result.get("nextCursor")
else:
# Defensive: tolerate providers that still return a raw list
models = result
next_cursor = None
if not isinstance(models, list):
models = []
if next_cursor is not None and not isinstance(next_cursor, str):
next_cursor = str(next_cursor)
estimated_total = None
if cursor is None:
get_count = getattr(metadata_provider, "get_creator_model_count", None)
if get_count is not None:
try:
estimated_total = await get_count(username)
except Exception: # best-effort only
estimated_total = None
if not isinstance(estimated_total, int):
estimated_total = None
lora_scanner = await self._service_registry.get_lora_scanner()
checkpoint_scanner = await self._service_registry.get_checkpoint_scanner()
@@ -2621,15 +2653,16 @@ class ModelLibraryHandler:
}
lora_type_aliases = {model_type.lower() for model_type in VALID_LORA_TYPES}
type_scanner_map: Dict[str, object | None] = {
type_scanner_map: Dict[str, Any] = {
**{alias: lora_scanner for alias in lora_type_aliases},
"checkpoint": checkpoint_scanner,
"textualinversion": embedding_scanner,
}
versions: list[dict] = []
versions: list[dict[str, Any]] = []
history_service = await self._get_download_history_service()
model_ids: list[int] = []
model_count = 0
for model in models:
try:
model_ids.append(int(model.get("id")))
@@ -2663,6 +2696,8 @@ class ModelLibraryHandler:
if model_type not in normalized_allowed_types:
continue
model_count += 1
scanner = type_scanner_map.get(model_type)
if scanner is None:
return web.json_response(
@@ -2676,6 +2711,8 @@ class ModelLibraryHandler:
tags_value = model.get("tags")
tags = tags_value if isinstance(tags_value, list) else []
model_id = model.get("id")
if model_id is None:
continue
try:
model_id_int = int(model_id)
except (TypeError, ValueError):
@@ -2691,6 +2728,8 @@ class ModelLibraryHandler:
continue
version_id = version.get("id")
if version_id is None:
continue
try:
version_id_int = int(version_id)
except (TypeError, ValueError):
@@ -2728,7 +2767,15 @@ class ModelLibraryHandler:
)
return web.json_response(
{"success": True, "username": username, "versions": versions}
{
"success": True,
"username": username,
"versions": versions,
"modelCount": model_count,
"nextCursor": next_cursor,
"hasMore": next_cursor is not None,
"estimatedTotal": estimated_total,
}
)
except Exception as exc: # pragma: no cover - defensive logging
logger.error("Failed to get Civitai user models: %s", exc, exc_info=True)
@@ -2744,7 +2791,7 @@ class MetadataArchiveHandler:
] = get_metadata_archive_manager,
settings_service=None,
metadata_provider_updater: Callable[
[], Awaitable[None]
[], Awaitable[Any]
] = update_metadata_providers,
) -> None:
self._metadata_archive_manager_factory = metadata_archive_manager_factory
@@ -2891,7 +2938,7 @@ class BackupHandler:
if request.content_type.startswith("multipart/"):
reader = await request.multipart()
field = await reader.next()
field: Any = await reader.next()
uploaded = False
while field is not None:
if getattr(field, "filename", None):
@@ -3510,7 +3557,7 @@ class NodeRegistryHandler:
except (TypeError, ValueError):
parsed_node_id = node_identifier
payload: dict = {
payload: dict[str, Any] = {
"id": parsed_node_id,
"value": value,
"mode": mode,
@@ -3634,7 +3681,7 @@ class NodeRegistryHandler:
except (TypeError, ValueError):
parsed_node_id = node_identifier
payload: dict = {
payload: dict[str, Any] = {
"id": parsed_node_id,
"value": value,
"mode": mode,
@@ -3701,8 +3748,8 @@ class MiscHandlerSet:
doctor: DoctorHandler,
example_workflows: ExampleWorkflowsHandler,
base_model: BaseModelHandlerSet,
hf_handler: HfHandler | None = None,
agent_handler: AgentHandler | None = None,
hf_handler: Any = None,
agent_handler: Any = None,
) -> None:
self.health = health
self.settings = settings
+122 -21
View File
@@ -71,7 +71,7 @@ class ModelPageView:
self._server_i18n = server_i18n
self._logger = logger
def _load_supporters(self) -> dict:
def _load_supporters(self) -> dict[str, Any]:
"""Load supporters data from JSON file."""
try:
current_file = os.path.abspath(__file__)
@@ -152,7 +152,7 @@ class ModelPageView:
self._template_env.filters["t"] = (
self._server_i18n.create_template_filter()
)
self._template_env._i18n_filter_added = True # type: ignore[attr-defined]
self._template_env._i18n_filter_added = True # pyright: ignore[reportAttributeAccessIssue]
from ...services.llm_service import PROVIDER_PRESETS
@@ -199,7 +199,7 @@ class ModelListingHandler:
self,
*,
service,
parse_specific_params: Callable[[web.Request], Dict],
parse_specific_params: Callable[[web.Request], Dict[str, Any]],
logger: logging.Logger,
) -> None:
self._service = service
@@ -287,7 +287,7 @@ class ModelListingHandler:
)
return web.json_response({"error": str(exc)}, status=500)
def _parse_common_params(self, request: web.Request) -> Dict:
def _parse_common_params(self, request: web.Request) -> Dict[str, Any]:
page = int(request.query.get("page", "1"))
page_size = min(int(request.query.get("page_size", "20")), 100)
sort_by = request.query.get("sort_by", "name")
@@ -394,12 +394,14 @@ class ModelListingHandler:
)
# 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):
civitai_model_id = None
# Keep as string — could be an HF group key (e.g. "hf:user/repo")
pass
return {
"page": page,
@@ -537,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
@@ -544,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,
@@ -566,7 +595,12 @@ class ModelManagementHandler:
{"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:
@@ -624,7 +658,7 @@ class ModelManagementHandler:
try:
reader = await request.multipart()
field = await reader.next()
field: Any = await reader.next()
if field is None or field.name != "preview_file":
raise ValueError("Expected 'preview_file' field")
content_type = field.headers.get("Content-Type", "image/png")
@@ -666,7 +700,7 @@ class ModelManagementHandler:
{
"success": True,
"preview_url": config.get_preview_static_url(
result["preview_path"]
str(result["preview_path"])
),
"preview_nsfw_level": result["preview_nsfw_level"],
}
@@ -747,7 +781,7 @@ class ModelManagementHandler:
result = await self._preview_service.replace_preview(
model_path=model_path,
preview_data=preview_data,
preview_data=preview_bytes,
content_type=content_type,
original_filename=original_filename,
nsfw_level=nsfw_level,
@@ -759,7 +793,7 @@ class ModelManagementHandler:
{
"success": True,
"preview_url": config.get_preview_static_url(
result["preview_path"]
str(result["preview_path"])
),
"preview_nsfw_level": result["preview_nsfw_level"],
}
@@ -1454,8 +1488,73 @@ class ModelQueryHandler:
search = request.query.get("search", "").strip()
limit = min(int(request.query.get("limit", "15")), 100)
offset = max(0, int(request.query.get("offset", "0")))
folder = request.query.get("folder")
recursive = request.query.get("recursive", "true").lower() == "true"
base_models = list(request.query.getall("base_model", []))
model_types = list(request.query.getall("model_type", []))
tag_filters: Dict[str, str] = {}
for tag in request.query.getall("tag_include", []):
if tag:
tag_filters[tag] = "include"
for tag in request.query.getall("tag_exclude", []):
if tag:
tag_filters[tag] = "exclude"
auto_tag_filters: Dict[str, str] = {}
for tag in request.query.getall("auto_tag_include", []):
if tag:
auto_tag_filters[tag] = "include"
for tag in request.query.getall("auto_tag_exclude", []):
if tag:
auto_tag_filters[tag] = "exclude"
tag_logic = request.query.get("tag_logic", "any").lower()
if tag_logic not in ("any", "all"):
tag_logic = "any"
credit_required = request.query.get("credit_required")
if credit_required is not None:
credit_required = credit_required.lower() not in ("false", "0", "")
allow_selling_generated_content = request.query.get(
"allow_selling_generated_content"
)
if allow_selling_generated_content is not None:
allow_selling_generated_content = (
allow_selling_generated_content.lower() not in ("false", "0", "")
)
# The presence of the recursive param (always sent by the loras
# widget when filter mode is on) signals that the filter pipeline
# must run even when no concrete filter is set, so global settings
# like show_only_sfw stay consistent with the list endpoint.
apply_filters = (
"recursive" in request.query
or folder is not None
or bool(base_models)
or bool(model_types)
or bool(tag_filters)
or bool(auto_tag_filters)
or credit_required is not None
or allow_selling_generated_content is not None
)
matching_paths = await self._service.search_relative_paths(
search, limit, offset
search,
limit,
offset,
folder=folder,
recursive=recursive,
base_models=base_models,
model_types=model_types,
tags=tag_filters,
auto_tags=auto_tag_filters,
tag_logic=tag_logic,
credit_required=credit_required,
allow_selling_generated_content=allow_selling_generated_content,
apply_filters=apply_filters,
)
return web.json_response(
{"success": True, "relative_paths": matching_paths}
@@ -1961,7 +2060,7 @@ class ModelCivitaiHandler:
settings_service: SettingsManager,
ws_manager: WebSocketManager,
logger: logging.Logger,
metadata_provider_factory: Callable[[], Awaitable],
metadata_provider_factory: Callable[[], Awaitable[Any]],
validate_model_type: Callable[[str], bool],
expected_model_types: Callable[[], str],
find_model_file: Callable[
@@ -2026,7 +2125,7 @@ class ModelCivitaiHandler:
downloaded_version_ids = set(
await history_service.get_downloaded_version_ids(
self._service.model_type,
model_id,
int(model_id),
)
)
except Exception as exc: # pragma: no cover - defensive logging
@@ -2303,8 +2402,8 @@ class ModelUpdateHandler:
self._logger.error("Failed to fetch license info: %s", exc, exc_info=True)
return web.json_response({"success": False, "error": str(exc)}, status=500)
updated: List[Dict[str, str]] = []
errors: List[Dict[str, str]] = []
updated: List[Dict[str, Any]] = []
errors: List[Dict[str, Any]] = []
for model_id in model_ids:
license_payload = license_map.get(model_id)
if not license_payload:
@@ -2317,6 +2416,7 @@ class ModelUpdateHandler:
model_section = civitai_section.get("model")
if not isinstance(model_section, Mapping):
model_section = {}
model_section = dict(model_section)
model_section.update(resolved_payload)
civitai_section["model"] = model_section
metadata_payload["civitai"] = civitai_section
@@ -2332,7 +2432,7 @@ class ModelUpdateHandler:
)
errors.append({"filePath": metadata_path, "error": str(exc)})
response_payload = {"success": True, "updated": updated}
response_payload: Dict[str, Any] = {"success": True, "updated": updated}
missing_model_ids = [mid for mid in model_ids if mid not in license_map]
if missing_model_ids:
response_payload["missingModelIds"] = missing_model_ids
@@ -2681,6 +2781,7 @@ class ModelUpdateHandler:
civitai_payload = metadata_payload.get("civitai")
if not isinstance(civitai_payload, Mapping):
civitai_payload = {}
civitai_payload = dict(civitai_payload)
model_payload = civitai_payload.get("model")
if not isinstance(model_payload, Mapping):
@@ -2725,7 +2826,7 @@ class ModelUpdateHandler:
return aggregated
def _extract_target_model_ids(self, payload: Dict) -> Optional[List[int]]:
def _extract_target_model_ids(self, payload: Dict[str, Any]) -> Optional[List[int]]:
if not isinstance(payload, Mapping):
return None
@@ -2753,7 +2854,7 @@ class ModelUpdateHandler:
return {}
to_dict = getattr(metadata, "to_dict", None)
if callable(to_dict):
if to_dict:
try:
return to_dict()
except Exception:
@@ -2764,7 +2865,7 @@ class ModelUpdateHandler:
return {}
async def _read_json(self, request: web.Request) -> Dict:
async def _read_json(self, request: web.Request) -> Dict[str, Any]:
if not request.can_read_body:
return {}
try:
@@ -2796,7 +2897,7 @@ class ModelUpdateHandler:
record,
*,
version_context: Optional[Dict[int, Dict[str, Any]]] = None,
) -> Dict:
) -> Dict[str, Any]:
context = version_context or {}
# Check user setting for hiding early access versions
hide_early_access = False
@@ -2825,7 +2926,7 @@ class ModelUpdateHandler:
@staticmethod
def _serialize_version(
version, context: Optional[Dict[str, Any]]
) -> Dict:
) -> Dict[str, Any]:
context = context or {}
preview_override = context.get("preview_override")
preview_url = (
+148 -53
View File
@@ -10,7 +10,7 @@ import asyncio
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Tuple
from aiohttp import web
@@ -44,6 +44,22 @@ EnsureDependenciesCallable = Callable[[], Awaitable[None]]
RecipeScannerGetter = Callable[[], Any]
CivitaiClientGetter = Callable[[], Any]
# Cap concurrent preview-dimension reads across requests. With a cold LRU
# cache one page can touch up to page_size image files; 16 balances SSD and
# HDD throughput without starving the event loop.
_DIMS_READ_SEMAPHORE = asyncio.Semaphore(16)
async def _read_preview_dims(path: str) -> Optional[Tuple[int, int]]:
"""Read preview dimensions off the event loop under the concurrency cap.
PIL I/O runs in a worker thread so it never blocks the event loop, and the
semaphore bounds how many files are opened at once even when many list
requests land together.
"""
async with _DIMS_READ_SEMAPHORE:
return await asyncio.to_thread(ExifUtils.get_image_dimensions, path)
@dataclass(frozen=True)
class RecipeHandlerSet:
@@ -246,7 +262,8 @@ class RecipeListingHandler:
recursive=recursive,
)
for item in result.get("items", []):
items = result.get("items", [])
for item in items:
file_path = item.get("file_path")
if file_path:
item["file_url"] = self.format_recipe_file_url(file_path)
@@ -255,6 +272,26 @@ class RecipeListingHandler:
item.setdefault("loras", [])
item.setdefault("base_model", "")
# Batch preview dimension reads with asyncio.gather. The previous
# loop awaited asyncio.to_thread once per item, so a page_size=100
# request submitted 100 sequential thread calls (50-300ms cold-page
# latency). gather runs them concurrently while the semaphore caps
# disk opens; dimensions stay omitted (not null) when a preview has
# no readable size (video, missing file).
to_read = [
(i, item.get("file_path"))
for i, item in enumerate(items)
if item.get("file_path")
]
if to_read:
dims_list = await asyncio.gather(
*(_read_preview_dims(path) for _, path in to_read)
)
for (idx, _), dims in zip(to_read, dims_list):
if dims:
item = items[idx]
item["width"], item["height"] = dims
return web.json_response(result)
except Exception as exc:
self._logger.error("Error retrieving recipes: %s", exc, exc_info=True)
@@ -1045,10 +1082,10 @@ class RecipeManagementHandler:
*,
image_url: str,
name: str,
lora_entries: list,
checkpoint_entry: dict,
gen_params_request: dict,
tags: list,
lora_entries: list[Any],
checkpoint_entry: Dict[str, Any] | None,
gen_params_request: Dict[str, Any] | None,
tags: list[Any],
base_model: str,
source_path: str,
) -> web.Response:
@@ -1081,6 +1118,12 @@ class RecipeManagementHandler:
_original_image_url,
) = await self._download_remote_media(image_url)
# Build a version-cached map of local model hashes to cache items so
# CivitaiApiMetadataParser can skip CivitAI API calls for models that
# exist on disk. Built once and shared by every parse pass below.
local_cache = await recipe_scanner.build_local_hash_cache()
from ...recipes.parsers.civitai_image import CivitaiApiMetadataParser
# Extract embedded EXIF metadata (offloaded to thread pool in this call)
embedded_gen_params = {}
parsed_embedded = None
@@ -1102,9 +1145,16 @@ class RecipeManagementHandler:
)
)
if parser:
parsed_embedded = await parser.parse_metadata(
raw_embedded, recipe_scanner=recipe_scanner
)
if isinstance(parser, CivitaiApiMetadataParser):
parsed_embedded = await parser.parse_metadata(
raw_embedded,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
)
else:
parsed_embedded = await parser.parse_metadata(
raw_embedded, recipe_scanner=recipe_scanner
)
if parsed_embedded and "gen_params" in parsed_embedded:
embedded_gen_params = parsed_embedded["gen_params"]
else:
@@ -1135,9 +1185,16 @@ class RecipeManagementHandler:
civitai_inner_meta
)
if parser:
civitai_parsed = await parser.parse_metadata(
civitai_inner_meta, recipe_scanner=recipe_scanner
)
if isinstance(parser, CivitaiApiMetadataParser):
civitai_parsed = await parser.parse_metadata(
civitai_inner_meta,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
)
else:
civitai_parsed = await parser.parse_metadata(
civitai_inner_meta, recipe_scanner=recipe_scanner
)
if civitai_parsed and "gen_params" in civitai_parsed:
# Merge: API gen_params override EXIF at field level,
# EXIF fills in fields the API doesn't have.
@@ -1641,7 +1698,7 @@ class RecipeManagementHandler:
if not provider:
return ""
version_info = await provider.get_model_version_info(version_id)
version_info = await provider.get_model_version_info(str(version_id))
if isinstance(version_info, tuple):
version_info = version_info[0]
@@ -1761,6 +1818,12 @@ class RecipeManagementHandler:
await self._download_remote_media(image_url)
)
# Build a version-cached map of local model hashes to cache items so
# CivitaiApiMetadataParser can skip CivitAI API calls for models that
# exist on disk. Built once and shared by every parse pass below.
local_cache = await recipe_scanner.build_local_hash_cache()
from ...recipes.parsers.civitai_image import CivitaiApiMetadataParser
# Extract embedded EXIF metadata
embedded_gen_params = {}
parsed_embedded = None
@@ -1782,9 +1845,16 @@ class RecipeManagementHandler:
)
)
if parser:
parsed_embedded = await parser.parse_metadata(
raw_embedded, recipe_scanner=recipe_scanner
)
if isinstance(parser, CivitaiApiMetadataParser):
parsed_embedded = await parser.parse_metadata(
raw_embedded,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
)
else:
parsed_embedded = await parser.parse_metadata(
raw_embedded, recipe_scanner=recipe_scanner
)
if parsed_embedded and "gen_params" in parsed_embedded:
embedded_gen_params = parsed_embedded["gen_params"]
finally:
@@ -1822,9 +1892,16 @@ class RecipeManagementHandler:
)
)
if parser:
parsed_embedded = await parser.parse_metadata(
raw_orig, recipe_scanner=recipe_scanner
)
if isinstance(parser, CivitaiApiMetadataParser):
parsed_embedded = await parser.parse_metadata(
raw_orig,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
)
else:
parsed_embedded = await parser.parse_metadata(
raw_orig, recipe_scanner=recipe_scanner
)
if (
parsed_embedded
and "gen_params" in parsed_embedded
@@ -1858,9 +1935,16 @@ class RecipeManagementHandler:
civitai_inner_meta
)
if parser:
civitai_parsed = await parser.parse_metadata(
civitai_inner_meta, recipe_scanner=recipe_scanner
)
if isinstance(parser, CivitaiApiMetadataParser):
civitai_parsed = await parser.parse_metadata(
civitai_inner_meta,
recipe_scanner=recipe_scanner,
local_cache=local_cache,
)
else:
civitai_parsed = await parser.parse_metadata(
civitai_inner_meta, recipe_scanner=recipe_scanner
)
if civitai_parsed and "gen_params" in civitai_parsed:
# Merge: API gen_params override EXIF at field level,
# EXIF fills in fields the API doesn't have.
@@ -2072,33 +2156,44 @@ class RecipeManagementHandler:
parsed_input = {**image_data, **inner_meta}
parsed_input.pop("meta", None)
# Build a local cache of {hash cache_item} so the parser can
# skip CivitAI API calls for models that exist on disk.
local_cache: Dict[str, Dict[str, Any]] = {}
lora_scanner = getattr(recipe_scanner, "_lora_scanner", None)
if lora_scanner and model_hash:
try:
parent_cache_data = await lora_scanner.get_cached_data()
for item in getattr(parent_cache_data, "raw_data", []):
if item.get("sha256", "").lower() == model_hash.lower():
local_cache[model_hash.lower()] = item
# Compute AutoV3 so the parser can also match on
# that hash type (CivitAI metadata resources use
# AutoV3).
file_path = item.get("file_path")
if file_path and os.path.exists(file_path):
try:
from ...utils.file_utils import (
calculate_autov3,
)
autov3 = calculate_autov3(file_path)
if autov3:
local_cache[autov3.lower()] = item
except Exception:
pass
break
except Exception:
pass
# Build the shared local hash cache so the parser can skip CivitAI
# API calls for models that exist on disk.
local_cache: Dict[str, Dict[str, Any]] = (
await recipe_scanner.build_local_hash_cache()
)
# Bounded supplement for un-backfilled parents. The shared builder
# never computes autov3; when the parent model exists on disk but
# its cached entry has no stored AutoV3, compute it for that single
# file and register the AutoV3 key so the parser can also match on
# that hash type (CivitAI metadata resources use AutoV3). This runs
# whenever the parent is found with an empty autov3, independent of
# whether the sha256 key is already present in the shared cache.
if model_hash:
lora_scanner = getattr(recipe_scanner, "_lora_scanner", None)
if lora_scanner:
try:
parent_cache_data = await lora_scanner.get_cached_data()
for item in getattr(parent_cache_data, "raw_data", []):
if item.get("sha256", "").lower() == model_hash.lower():
autov3 = (item.get("autov3") or "").lower()
if not autov3:
file_path = item.get("file_path")
if file_path and os.path.exists(file_path):
try:
from ...utils.file_utils import (
calculate_autov3,
)
autov3 = (
calculate_autov3(file_path) or ""
).lower()
except Exception:
pass
if autov3:
local_cache[autov3] = item
break
except Exception:
pass
parser = self._analysis_service._recipe_parser_factory.create_parser(
parsed_input
@@ -2130,10 +2225,10 @@ class RecipeManagementHandler:
parent_model_id: int | None = None
parent_version_name: str | None = None
parent_model_name: str | None = None
# Prefer sha256 key; fall back to any cached entry.
# Resolve the parent strictly by its sha256 key. There is no
# arbitrary fallback: with a full-library cache, picking any entry
# would corrupt the isDeleted reconciliation below.
parent_item = local_cache.get(model_hash.lower()) if model_hash else None
if parent_item is None and local_cache:
parent_item = next(iter(local_cache.values()))
if parent_item:
civ = parent_item.get("civitai") or {}
if isinstance(civ, dict):
@@ -2349,7 +2444,7 @@ class RecipeAnalysisHandler:
content_type = request.headers.get("Content-Type", "")
if "multipart/form-data" in content_type:
reader = await request.multipart()
field = await reader.next()
field: Any = await reader.next()
if field is None or field.name != "image":
raise RecipeValidationError("No image field found")
image_chunks = bytearray()
+6 -71
View File
@@ -1,8 +1,8 @@
import asyncio
import logging
from aiohttp import web
from typing import Dict
from server import PromptServer # type: ignore
from typing import Any, Dict
from server import PromptServer # pyright: ignore[reportMissingImports]
from .base_model_routes import BaseModelRoutes
from .model_route_registrar import ModelRouteRegistrar
@@ -31,13 +31,13 @@ class LoraRoutes(BaseModelRoutes):
# Attach service dependencies
self.attach_service(self.service)
def setup_routes(self, app: web.Application):
def setup_routes(self, app: web.Application, prefix: str = "loras"):
"""Setup LoRA routes"""
# Schedule service initialization on app startup
app.on_startup.append(lambda _: self.initialize_services())
# Setup common routes with 'loras' prefix (includes page route)
super().setup_routes(app, "loras")
super().setup_routes(app, prefix)
def setup_specific_routes(self, registrar: ModelRouteRegistrar, prefix: str):
"""Setup LoRA-specific routes"""
@@ -73,7 +73,7 @@ class LoraRoutes(BaseModelRoutes):
"POST", "/api/lm/{prefix}/get_trigger_words", prefix, self.get_trigger_words
)
def _parse_specific_params(self, request: web.Request) -> Dict:
def _parse_specific_params(self, request: web.Request) -> Dict[str, Any]:
"""Parse LoRA-specific parameters"""
params = {}
@@ -119,25 +119,6 @@ class LoraRoutes(BaseModelRoutes):
logger.error(f"Error getting letter counts: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
async def get_lora_notes(self, request: web.Request) -> web.Response:
"""Get notes for a specific LoRA file"""
try:
lora_name = request.query.get("name")
if not lora_name:
return web.Response(text="Lora file name is required", status=400)
notes = await self.service.get_lora_notes(lora_name)
if notes is not None:
return web.json_response({"success": True, "notes": notes})
else:
return web.json_response(
{"success": False, "error": "LoRA not found in cache"}, status=404
)
except Exception as e:
logger.error(f"Error getting lora notes: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
async def get_lora_trigger_words(self, request: web.Request) -> web.Response:
"""Get trigger words for a specific LoRA file"""
try:
@@ -168,52 +149,6 @@ class LoraRoutes(BaseModelRoutes):
logger.error(f"Error getting lora usage tips by path: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
async def get_lora_preview_url(self, request: web.Request) -> web.Response:
"""Get the static preview URL for a LoRA file"""
try:
lora_name = request.query.get("name")
if not lora_name:
return web.Response(text="Lora file name is required", status=400)
preview_url = await self.service.get_lora_preview_url(lora_name)
if preview_url:
return web.json_response({"success": True, "preview_url": preview_url})
else:
return web.json_response(
{
"success": False,
"error": "No preview URL found for the specified lora",
},
status=404,
)
except Exception as e:
logger.error(f"Error getting lora preview URL: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
async def get_lora_civitai_url(self, request: web.Request) -> web.Response:
"""Get the Civitai URL for a LoRA file"""
try:
lora_name = request.query.get("name")
if not lora_name:
return web.Response(text="Lora file name is required", status=400)
result = await self.service.get_lora_civitai_url(lora_name)
if result["civitai_url"]:
return web.json_response({"success": True, **result})
else:
return web.json_response(
{
"success": False,
"error": "No Civitai data found for the specified lora",
},
status=404,
)
except Exception as e:
logger.error(f"Error getting lora Civitai URL: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
async def get_random_loras(self, request: web.Request) -> web.Response:
"""Get random LoRAs based on filters and strength ranges"""
try:
@@ -337,7 +272,7 @@ class LoraRoutes(BaseModelRoutes):
graph_identifier = entry.get("graph_id")
try:
parsed_node_id = int(node_identifier)
parsed_node_id = int(node_identifier) # pyright: ignore[reportArgumentType]
except (TypeError, ValueError):
parsed_node_id = node_identifier
+2 -2
View File
@@ -5,7 +5,7 @@ miscellaneous endpoints share a consistent registration flow.
"""
from dataclasses import dataclass
from typing import Callable, Iterable, Mapping
from typing import Any, Callable, Iterable, Mapping
from aiohttp import web
@@ -147,7 +147,7 @@ class MiscRouteRegistrar:
handler_lookup[definition.handler_name],
)
def _bind(self, method: str, path: str, handler: Callable) -> None:
def _bind(self, method: str, path: str, handler: Callable[..., Any]) -> None:
add_method_name = self._METHOD_MAP[method.upper()]
add_method = getattr(self._app.router, add_method_name)
add_method(path, handler)
+1 -1
View File
@@ -7,7 +7,7 @@ import os
from typing import Awaitable, Callable, Mapping
from aiohttp import web
from server import PromptServer # type: ignore
from server import PromptServer # pyright: ignore[reportMissingImports]
from ..services.metadata_service import (
get_metadata_archive_manager,
+4 -4
View File
@@ -3,7 +3,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable, Iterable, Mapping
from typing import Any, Callable, Iterable, Mapping
from aiohttp import web
@@ -174,15 +174,15 @@ class ModelRouteRegistrar:
handler_lookup[definition.handler_name],
)
def add_route(self, method: str, path: str, handler: Callable) -> None:
def add_route(self, method: str, path: str, handler: Callable[..., Any]) -> None:
self._bind_route(method, path, handler)
def add_prefixed_route(
self, method: str, path_template: str, prefix: str, handler: Callable
self, method: str, path_template: str, prefix: str, handler: Callable[..., Any]
) -> None:
self._bind_route(method, path_template.replace("{prefix}", prefix), handler)
def _bind_route(self, method: str, path: str, handler: Callable) -> None:
def _bind_route(self, method: str, path: str, handler: Callable[..., Any]) -> None:
add_method_name = self._METHOD_MAP[method.upper()]
add_method = getattr(self._app.router, add_method_name)
add_method(path, handler)
+2 -2
View File
@@ -3,7 +3,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable, Mapping
from typing import Any, Callable, Mapping
from aiohttp import web
@@ -105,7 +105,7 @@ class RecipeRouteRegistrar:
handler = handler_lookup[definition.handler_name]
self._bind_route(definition.method, definition.path, handler)
def _bind_route(self, method: str, path: str, handler: Callable) -> None:
def _bind_route(self, method: str, path: str, handler: Callable[..., Any]) -> None:
add_method_name = self._METHOD_MAP[method.upper()]
add_method = getattr(self._app.router, add_method_name)
add_method(path, handler)
+11 -10
View File
@@ -40,10 +40,11 @@ class StatsRoutes:
"""Route handlers for Statistics page and API endpoints"""
def __init__(self):
self.lora_scanner = None
self.checkpoint_scanner = None
self.embedding_scanner = None
self.usage_stats = None
self.lora_scanner: Any = None
self.checkpoint_scanner: Any = None
self.embedding_scanner: Any = None
self.usage_stats: Any = None
self._i18n_filter_added = False
self.template_env = jinja2.Environment(
loader=jinja2.FileSystemLoader(config.templates_path),
autoescape=True
@@ -95,9 +96,9 @@ class StatsRoutes:
server_i18n.set_locale(user_language)
# 为模板环境添加i18n过滤器
if not hasattr(self.template_env, '_i18n_filter_added'):
if not self._i18n_filter_added:
self.template_env.filters['t'] = server_i18n.create_template_filter()
self.template_env._i18n_filter_added = True
self._i18n_filter_added = True
template = self.template_env.get_template('statistics.html')
rendered = template.render(
@@ -549,7 +550,7 @@ class StatsRoutes:
'error': str(e)
}, status=500)
def _count_unused_models(self, models: List[Dict], usage_data: Dict) -> int:
def _count_unused_models(self, models: List[Dict[str, Any]], usage_data: Dict[str, Any]) -> int:
"""Count models that have never been used"""
used_hashes = set(usage_data.keys())
unused_count = 0
@@ -560,7 +561,7 @@ class StatsRoutes:
return unused_count
def _get_top_used_models(self, usage_data: Dict, model_map: Dict, limit: int) -> List[Dict]:
def _get_top_used_models(self, usage_data: Dict[str, Any], model_map: Dict[str, Any], limit: int) -> List[Dict[str, Any]]:
"""Get top used models with their metadata"""
sorted_usage = sorted(usage_data.items(), key=lambda x: x[1].get('total', 0), reverse=True)
@@ -578,7 +579,7 @@ class StatsRoutes:
return top_models
def _get_usage_timeline(self, usage_data: Dict, days: int) -> List[Dict]:
def _get_usage_timeline(self, usage_data: Dict[str, Any], days: int) -> List[Dict[str, Any]]:
"""Get usage timeline for the past N days"""
timeline = []
today = datetime.now()
@@ -614,7 +615,7 @@ class StatsRoutes:
return list(reversed(timeline)) # Oldest to newest
def _format_size(self, size_bytes: int) -> str:
def _format_size(self, size_bytes: float) -> str:
"""Format file size in human readable format"""
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
if size_bytes < 1024.0:
+324 -53
View File
@@ -6,7 +6,7 @@ import shutil
import tempfile
import asyncio
from aiohttp import web, ClientError
from typing import Dict, List
from typing import Any, Dict, List, cast
from ..utils.settings_paths import ensure_settings_file
from ..services.downloader import get_downloader
@@ -38,6 +38,84 @@ def _clean_excludes() -> List[str]:
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"""
@@ -47,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):
@@ -65,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:
@@ -81,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,
@@ -88,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:
@@ -126,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:
@@ -156,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({
@@ -190,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]:
"""
@@ -212,9 +467,10 @@ class UpdateRoutes:
if not success:
logger.error(f"Failed to fetch release info: {data}")
return False, ""
zip_url = data.get("zipball_url")
version = data.get("tag_name", "unknown")
release_payload = cast(dict[str, Any], data)
zip_url = release_payload.get("zipball_url", "")
version = release_payload.get("tag_name", "unknown")
# Download ZIP to temporary file
with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp_zip:
@@ -244,8 +500,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:
@@ -255,7 +510,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
@@ -272,7 +527,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)
@@ -295,7 +550,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):
@@ -308,41 +564,56 @@ 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
commit_message = data.get('commit', {}).get('message', '')
# Format as "main-{short_hash}"
logger.warning("Failed to fetch GitHub commit: %s", data)
return "main", [], 0, ""
commit_payload = cast(dict[str, Any], data)
commit_sha = commit_payload.get('sha', '')[:7]
commit_message = commit_payload.get('commit', {}).get('message', '')
commit_date = commit_payload.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:
compare_payload = cast(dict[str, Any], c_data)
if compare_payload.get('status') in ('ahead', 'diverged'):
behind_by = compare_payload.get('ahead_by', 0)
else:
behind_by = compare_payload.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:
@@ -438,7 +709,7 @@ class UpdateRoutes:
logger.info(f"Successfully updated to {new_version}")
return True, new_version
except git.exc.GitError as e:
except git.exc.GitError as e: # pyright: ignore[reportAttributeAccessIssue]
logger.error(f"Git error during update: {e}")
return False, ""
except Exception as e:
@@ -499,7 +770,7 @@ class UpdateRoutes:
return git_info
@staticmethod
async def _get_remote_version() -> tuple[str, List[str], List[Dict]]:
async def _get_remote_version() -> tuple[str, List[str], List[Dict[str, Any]]]:
"""
Fetch remote version from GitHub
Returns:
@@ -521,7 +792,7 @@ class UpdateRoutes:
# Parse releases
releases = []
for i, release in enumerate(data):
for i, release in enumerate(cast(list[dict[str, Any]], data)):
version = release.get('tag_name', '')
if not version.startswith('v'):
version = f"v{version}"
+1 -1
View File
@@ -117,7 +117,7 @@ def _render_prompt(template: str, variables: Dict[str, Any]) -> str:
Uses simple regex substitution no Jinja2 dependency needed.
"""
def replace(match: re.Match) -> str:
def replace(match: re.Match[str]) -> str:
key = match.group(1).strip()
value = variables.get(key, "")
if isinstance(value, (dict, list)):
+1 -1
View File
@@ -295,7 +295,7 @@ class PostProcessor:
normalises every tag to lowercase for case-insensitive dedup.
"""
merged: List[str] = []
seen: set = set()
seen: set[str] = set()
for tag in list(existing) + list(new):
t = tag.strip().lower()
if t and t not in seen:
+1 -1
View File
@@ -49,7 +49,7 @@ _FRONTMATTER_RE = re.compile(
)
def _parse_skill_file(path: Path) -> tuple[dict, str]:
def _parse_skill_file(path: Path) -> tuple[dict[str, Any], str]:
"""Read a prompt definition file (``prompt.md`` or legacy ``SKILL.md``) and
return (frontmatter_dict, body_text).
@@ -9,7 +9,7 @@ from __future__ import annotations
import html as html_module
import re
from typing import List, Tuple
from typing import Any, List, Tuple
_REPO_URL_PATTERN = re.compile(r"https?://huggingface\.co/([^/]+/[^/]+)")
@@ -18,10 +18,10 @@ _REPO_URL_PATTERN = re.compile(r"https?://huggingface\.co/([^/]+/[^/]+)")
def extract_simple_markdown_images(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
existing_urls: set[str] | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
) -> list[dict[str, Any]]:
"""Extract standalone markdown images from the README body.
Matches ``![alt](url)`` on lines that are NOT part of a markdown table
@@ -36,8 +36,8 @@ def extract_simple_markdown_images(
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
images: list[dict[str, Any]] = []
seen_urls: set[str] = set(existing_urls) if existing_urls else set()
# Collect lines that are NOT inside fenced code blocks
lines = markdown_text.split("\n")
@@ -86,10 +86,10 @@ def extract_simple_markdown_images(
def extract_html_img_tags(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
existing_urls: set[str] | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
) -> list[dict[str, Any]]:
"""Extract image URLs from HTML ``<img src=\"...\">`` tags in the README.
Many HF collection repos (e.g. ``deadman44/Z-Image_LoRA``) use raw HTML
@@ -103,8 +103,8 @@ def extract_html_img_tags(
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
images: list[dict[str, Any]] = []
seen_urls: set[str] = set(existing_urls) if existing_urls else set()
for m in re.finditer(
r'<img\s[^>]*src=\"([^\"]+)\"',
@@ -175,7 +175,7 @@ def extract_gallery_images(
repo: str,
default_width: int = 512,
default_height: int = 512,
) -> List[dict]:
) -> List[dict[str, Any]]:
"""Extract widget/gallery images from the YAML frontmatter of a HF README.
Args:
@@ -196,7 +196,7 @@ def extract_gallery_images(
if not frontmatter:
return []
images: List[dict] = []
images: List[dict[str, Any]] = []
base_url = f"https://huggingface.co/{repo}/resolve/main"
w = default_width or 512
h = default_height or 512
@@ -258,7 +258,7 @@ def extract_gallery_images(
text = raw_text
if url:
image: dict = {
image: dict[str, Any] = {
"url": url,
"type": "image",
"nsfwLevel": 0,
@@ -276,10 +276,10 @@ def extract_gallery_images(
def extract_gallery_table_images(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
existing_urls: set[str] | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
) -> list[dict[str, Any]]:
"""Extract images from ``| Preview | Prompt |`` markdown gallery tables.
Many HF READMEs include a sample-gallery table in the body (outside
@@ -295,8 +295,8 @@ def extract_gallery_table_images(
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
images: list[dict[str, Any]] = []
seen_urls: set[str] = set(existing_urls) if existing_urls else set()
lines = markdown_text.split("\n")
n = len(lines)
i = 0
@@ -514,7 +514,7 @@ def _strip_standalone_images(text: str) -> str:
URL was stripped entirely, making it impossible for the LLM to return
a ``preview_url`` for repos that use HTML ``<img>`` tags exclusively.
"""
def _img_to_md(match: re.Match) -> str:
def _img_to_md(match: re.Match[str]) -> str:
"""Convert an ``<img>`` tag to markdown image syntax ``![alt](src)``."""
tag = match.group(0)
src_m = re.search(r'src="([^"]+)"', tag) or re.search(r"src='([^']+)'", tag)
@@ -942,7 +942,7 @@ def _strip_badge_images(text: str) -> str:
"twitter", "colab", "gradio", "space",
)
def _should_remove(m: re.Match) -> str:
def _should_remove(m: re.Match[str]) -> str:
alt = (m.group(1) or "").lower()
for kw in badge_keywords:
if kw in alt:
+7 -3
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
from __future__ import annotations
import asyncio
@@ -23,7 +27,7 @@ logger = logging.getLogger(__name__)
def _try_certifi_ca_path() -> str | None:
"""Return the certifi CA bundle path if available, else None."""
try:
import certifi # type: ignore[import-untyped]
import certifi # pyright: ignore[reportMissingTypeStubs]
path = certifi.where()
if os.path.isfile(path):
@@ -84,7 +88,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
self._stderr_reader_task: Optional[asyncio.Task[Any]] = None
@property
def is_running(self) -> bool:
@@ -190,7 +194,7 @@ class Aria2Downloader:
download_id,
)
options: Dict[str, str] = {
options: Dict[str, Any] = {
"dir": save_dir,
"out": out_name,
"continue": "true",
+3 -3
View File
@@ -8,7 +8,7 @@ from filename, base_model, and CivitAI version name — no manual tagging requir
from __future__ import annotations
import re
from typing import Dict, List, Set
from typing import Any, Dict, List, Set
# ── Tag category definitions ──────────────────────────────────────────
# Each category maps a display label to a regex pattern.
@@ -52,7 +52,7 @@ AUTO_TAG_GROUPS = {
DEFAULT_ENABLED_GROUPS = {"mode", "video"}
def _collect_sources(model_data: Dict) -> List[str]:
def _collect_sources(model_data: Dict[str, Any]) -> List[str]:
"""Collect all text sources from model data for tag matching."""
sources: List[str] = []
@@ -73,7 +73,7 @@ def _collect_sources(model_data: Dict) -> List[str]:
return sources
def extract_auto_tags(model_data: Dict) -> List[str]:
def extract_auto_tags(model_data: Dict[str, Any]) -> List[str]:
"""Extract auto-detected tags from model metadata.
Uses a two-layer approach:
+144
View File
@@ -0,0 +1,144 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""Backfill the AutoV3 checked state for models loaded from a persisted snapshot.
The SQLite persistent cache predates the AutoV3 feature, so entries hydrated
from it have a NULL ``autov3`` column (the "not checked yet" state). This
service computes the embedded AutoV3 hash for each such model once per
process and persists it through the scanner's single write path
(:meth:`ModelScanner.update_autov3_for_model`), marking every visited row so a
subsequent run finds nothing left to do.
Three-state contract honored here:
- ``NULL`` (sqlite) / absent (dict) = not checked yet backfill computes it
- ``''`` (sqlite/dict) / JSON null = checked, no value available never recompute
- 12-char lowercase hex = value never recompute
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import threading
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING: # pragma: no cover - type-check only; runtime imports are local
from .model_scanner import ModelScanner
logger = logging.getLogger(__name__)
def _resolve_autov3(file_path: str) -> str:
"""Resolve the AutoV3 hash for a model file.
Prefers the Civitai AutoV3 reported for the file whose SHA256 matches
(the authoritative value for recipe matching); falls back to the embedded
safetensors header hash. Returns ``''`` when neither is available.
"""
try:
metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
if os.path.exists(metadata_path):
with open(metadata_path, "r", encoding="utf-8") as handle:
payload = json.load(handle)
if isinstance(payload, dict):
from ..utils.models import autov3_from_civitai_files # local import avoids cycles
sha256 = (payload.get("sha256") or "").lower()
civitai_autov3 = autov3_from_civitai_files(payload.get("civitai"), sha256)
if civitai_autov3:
return civitai_autov3
except Exception:
pass
from ..utils.file_utils import calculate_autov3 # local import avoids cycles
return calculate_autov3(file_path) or ""
class Autov3BackfillService:
"""Compute and persist AutoV3 hashes for models missing a checked state."""
_instance: Optional["Autov3BackfillService"] = None
_instance_lock = threading.Lock()
def __init__(self) -> None:
# Re-entrancy guard per model type: scanners for different model types
# initialize concurrently (lora_manager.py), so a global guard would
# silently skip every type but the first to start. Each model type
# runs its own backfill; a duplicate trigger for the same type no-ops.
self._running_types: set[str] = set()
@classmethod
def get_instance(cls) -> "Autov3BackfillService":
"""Return the process-wide singleton instance."""
if cls._instance is None:
with cls._instance_lock:
if cls._instance is None:
cls._instance = cls()
return cls._instance
async def backfill(self, scanner: "ModelScanner") -> int:
"""Compute AutoV3 for every un-checked model of ``scanner.model_type``.
Each candidate file is read once via :func:`~py.utils.file_utils.calculate_autov3`
(cheap: safetensors header only) and the result is persisted through
``scanner.update_autov3_for_model``. Files that no longer exist on
disk are skipped they are intentionally NOT marked, because scanner
cleanup removes the stale row later.
Returns:
The number of models successfully updated. Never raises; on any
failure a warning is logged and ``0`` is returned. A duplicate
trigger for a model type that is already being backfilled returns
``0`` immediately; different model types run concurrently.
"""
model_type = scanner.model_type
if model_type in self._running_types:
return 0
self._running_types.add(model_type)
try:
# Local imports avoid import cycles at module load time.
from .persistent_model_cache import get_persistent_cache
from ..utils.file_utils import calculate_autov3
persistent = getattr(scanner, "_persistent_cache", None) or get_persistent_cache()
paths = persistent.get_models_missing_autov3(model_type)
loop = asyncio.get_running_loop()
count = 0
for path in paths:
# A file that no longer exists must not be marked; scanner
# cleanup removes the stale row later. The existence check and
# hash resolution run in the executor so the loop stays
# responsive to API requests while the backfill iterates a
# large library.
if not await loop.run_in_executor(None, os.path.exists, path):
continue
autov3 = await loop.run_in_executor(None, _resolve_autov3, path)
if await scanner.update_autov3_for_model(model_type, path, autov3):
count += 1
if paths:
logger.info(
"AutoV3 backfill: updated %d/%d models for %s",
count,
len(paths),
model_type,
)
else:
# Steady state after the first run: nothing left to backfill.
logger.debug("AutoV3 backfill: nothing to process for %s", model_type)
return count
except Exception as exc:
logger.warning(
"AutoV3 backfill failed for %s: %s",
getattr(scanner, "model_type", "?"),
exc,
)
return 0
finally:
self._running_types.discard(model_type)
+4
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
from __future__ import annotations
import asyncio
+195 -73
View File
@@ -1,7 +1,8 @@
from abc import ABC, abstractmethod
import asyncio
import re
from typing import Any, Dict, List, Optional, Type, TYPE_CHECKING
import random
from typing import Any, Awaitable, Dict, List, Optional, Type, Union, TYPE_CHECKING, cast
import logging
import os
import time
@@ -69,24 +70,24 @@ class BaseModelService(ABC):
page: int,
page_size: int,
sort_by: str = "name",
folder: str = None,
folder_include: list = None,
folder_exclude: list = None,
search: str = None,
folder: str | None = None,
folder_include: list[str] | None = None,
folder_exclude: list[str] | None = None,
search: str | None = None,
fuzzy_search: bool = False,
base_models: list = None,
model_types: list = None,
base_models: list[str] | None = None,
model_types: list[str] | None = None,
tags: Optional[Dict[str, str]] = None,
auto_tags: Optional[Dict[str, str]] = None,
search_options: dict = None,
hash_filters: dict = None,
search_options: dict[str, Any] | None = None,
hash_filters: dict[str, Any] | None = None,
favorites_only: bool = False,
update_available_only: bool = False,
credit_required: Optional[bool] = None,
allow_selling_generated_content: Optional[bool] = None,
tag_logic: str = "any",
**kwargs,
) -> Dict:
) -> Dict[str, Any]:
"""Get paginated and filtered model data"""
overall_start = time.perf_counter()
@@ -109,12 +110,15 @@ class BaseModelService(ABC):
if civitai_model_id is not None:
sorted_data = [
item for item in sorted_data
if self._extract_model_id(item) == civitai_model_id
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 0,
key=lambda x: self._extract_version_id(x)
or x.get("modified", 0)
or 0,
reverse=True,
)
@@ -129,18 +133,21 @@ class BaseModelService(ABC):
ufs = self.settings.get("version_grouping", "same_base")
group_by_base = ufs == "same_base"
dedup_map = {} # (modelId [,base_model]) -> (item, version_id)
dedup_map = {} # (modelId [,base_model]) -> (item, version_or_modified)
version_counter = {} # same-key -> count
standalone = []
for item in sorted_data:
mid = self._extract_model_id(item)
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
vid = self._extract_version_id(item) or 0
# 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
@@ -171,19 +178,22 @@ class BaseModelService(ABC):
ufs = self.settings.get("version_grouping", "same_base")
group_by_base = ufs == "same_base"
model_groups: Dict[Any, List[Dict]] = {}
ungrouped_standalone: List[Dict] = []
model_groups: Dict[Any, List[Dict[str, Any]]] = {}
ungrouped_standalone: List[Dict[str, Any]] = []
for item in sorted_data:
mid = self._extract_model_id(item)
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
# 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 0,
key=lambda x: self._extract_version_id(x)
or x.get("modified", 0)
or 0,
reverse=True,
)
# Sort groups by version count
@@ -239,7 +249,7 @@ class BaseModelService(ABC):
filter_duration = time.perf_counter() - t1
post_filter_count = len(filtered_data)
annotated_for_filter: Optional[List[Dict]] = None
annotated_for_filter: Optional[List[Dict[str, Any]]] = None
t2 = time.perf_counter()
if update_available_only:
annotated_for_filter = await self._annotate_update_flags(filtered_data)
@@ -286,11 +296,11 @@ class BaseModelService(ABC):
page: int,
page_size: int,
sort_by: str = "name",
search: str = None,
search: str | None = None,
fuzzy_search: bool = False,
search_options: dict = None,
search_options: dict[str, Any] | None = None,
**kwargs,
) -> Dict:
) -> Dict[str, Any]:
"""Get paginated excluded model data."""
excluded_paths = list(self.scanner.get_excluded_models())
excluded_entries: List[Dict[str, Any]] = []
@@ -316,7 +326,7 @@ class BaseModelService(ABC):
]
persist_current_cache = getattr(self.scanner, "_persist_current_cache", None)
if callable(persist_current_cache):
await persist_current_cache()
await cast(Awaitable[Any], persist_current_cache())
excluded_entries = self._sort_entries(excluded_entries, sort_by)
@@ -381,6 +391,12 @@ class BaseModelService(ABC):
(item.get("model_name") or item.get("file_name") or "").lower(),
item.get("file_path", "").lower(),
)
elif key_name == "random":
# Seeded random shuffle: same seed -> same order (stable pagination)
rng = random.Random(sort_params.seed or "random")
result = list(data)
rng.shuffle(result)
return result
elif key_name == "size":
key_fn = lambda item: (
int(item.get("size", 0) or 0),
@@ -428,39 +444,50 @@ class BaseModelService(ABC):
return entry
async def _apply_hash_filters(
self, data: List[Dict], hash_filters: Dict
) -> List[Dict]:
"""Apply hash-based filtering"""
self, data: List[Dict[str, Any]], hash_filters: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""Apply hash-based filtering (SHA256 and AutoV3)."""
def matches_hash_set(item: Dict[str, Any], hash_set: set[str]) -> bool:
"""Check whether an item matches any hash in the set.
Compares the item's ``sha256`` field and its non-empty ``autov3``
field, both case-insensitively.
"""
if item.get("sha256", "").lower() in hash_set:
return True
autov3 = item.get("autov3", "")
return bool(autov3) and autov3.lower() in hash_set
single_hash = hash_filters.get("single_hash")
multiple_hashes = hash_filters.get("multiple_hashes")
if single_hash:
# Filter by single hash
single_hash = single_hash.lower()
# Filter by single hash (SHA256 or AutoV3)
return [
item for item in data if item.get("sha256", "").lower() == single_hash
item for item in data if matches_hash_set(item, {single_hash.lower()})
]
elif multiple_hashes:
# Filter by multiple hashes
hash_set = set(hash.lower() for hash in multiple_hashes)
return [item for item in data if item.get("sha256", "").lower() in hash_set]
# Filter by multiple hashes (SHA256 or AutoV3)
hash_set = {hash.lower() for hash in multiple_hashes}
return [item for item in data if matches_hash_set(item, hash_set)]
return data
async def _apply_common_filters(
self,
data: List[Dict],
folder: str = None,
folder_include: list = None,
folder_exclude: list = None,
base_models: list = None,
model_types: list = None,
data: List[Dict[str, Any]],
folder: str | None = None,
folder_include: list[str] | None = None,
folder_exclude: list[str] | None = None,
base_models: list[str] | None = None,
model_types: list[str] | None = None,
tags: Optional[Dict[str, str]] = None,
auto_tags: Optional[Dict[str, str]] = None,
favorites_only: bool = False,
search_options: dict = None,
search_options: dict[str, Any] | None = None,
tag_logic: str = "any",
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""Apply common filters that work across all model types"""
normalized_options = self.search_strategy.normalize_options(search_options)
criteria = FilterCriteria(
@@ -479,24 +506,24 @@ class BaseModelService(ABC):
async def _apply_search_filters(
self,
data: List[Dict],
data: List[Dict[str, Any]],
search: str,
fuzzy_search: bool,
search_options: dict,
) -> List[Dict]:
search_options: dict[str, Any] | None,
) -> List[Dict[str, Any]]:
"""Apply search filtering"""
normalized_options = self.search_strategy.normalize_options(search_options)
return self.search_strategy.apply(
data, search, normalized_options, fuzzy_search
)
async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]:
async def _apply_specific_filters(self, data: List[Dict[str, Any]], **kwargs) -> List[Dict[str, Any]]:
"""Apply model-specific filters - to be overridden by subclasses if needed"""
return data
async def _apply_credit_required_filter(
self, data: List[Dict], credit_required: bool
) -> List[Dict]:
self, data: List[Dict[str, Any]], credit_required: bool
) -> List[Dict[str, Any]]:
"""Apply credit required filtering based on license_flags.
Args:
@@ -526,8 +553,8 @@ class BaseModelService(ABC):
return filtered_data
async def _apply_allow_selling_filter(
self, data: List[Dict], allow_selling: bool
) -> List[Dict]:
self, data: List[Dict[str, Any]], allow_selling: bool
) -> List[Dict[str, Any]]:
"""Apply allow selling generated content filtering based on license_flags.
Args:
@@ -559,8 +586,8 @@ class BaseModelService(ABC):
async def _annotate_update_flags(
self,
items: List[Dict],
) -> List[Dict]:
items: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""Attach an update_available flag to each response item.
Items without a civitai model id default to False.
@@ -575,7 +602,7 @@ class BaseModelService(ABC):
item["update_available"] = False
return annotated
id_to_items: Dict[int, List[Dict]] = {}
id_to_items: Dict[int, List[Dict[str, Any]]] = {}
ordered_ids: List[int] = []
for item in annotated:
model_id = self._extract_model_id(item)
@@ -612,7 +639,7 @@ class BaseModelService(ABC):
record_method = getattr(self.update_service, "get_records_bulk", None)
if callable(record_method):
try:
records = await record_method(self.model_type, ordered_ids)
records = await cast(Awaitable[Any], record_method(self.model_type, ordered_ids))
resolved = {
model_id: record.has_update(hide_early_access=hide_early_access)
for model_id, record in records.items()
@@ -632,11 +659,11 @@ class BaseModelService(ABC):
bulk_method = getattr(self.update_service, "has_updates_bulk", None)
if callable(bulk_method):
try:
resolved = await bulk_method(
resolved = await cast(Awaitable[Any], bulk_method(
self.model_type,
ordered_ids,
hide_early_access=hide_early_access,
)
))
except Exception as exc:
logger.error(
"Failed to resolve update status in bulk for %s models (%s): %s",
@@ -698,7 +725,34 @@ class BaseModelService(ABC):
return annotated
@staticmethod
def _extract_model_id(item: Dict) -> Optional[int]:
def _extract_hf_group_key(item: Dict[str, Any]) -> 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[str, Any]) -> 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[str, Any]) -> Optional[int]:
civitai = item.get("civitai") if isinstance(item, dict) else None
if not isinstance(civitai, dict):
return None
@@ -711,7 +765,7 @@ class BaseModelService(ABC):
return None
@staticmethod
def _extract_version_id(item: Dict) -> Optional[int]:
def _extract_version_id(item: Dict[str, Any]) -> Optional[int]:
civitai = item.get("civitai") if isinstance(item, dict) else None
if not isinstance(civitai, dict):
return None
@@ -724,7 +778,7 @@ class BaseModelService(ABC):
return None
@staticmethod
def _extract_base_model(item: Dict) -> Optional[str]:
def _extract_base_model(item: Dict[str, Any]) -> Optional[str]:
value = item.get("base_model")
if value is None:
return None
@@ -776,7 +830,7 @@ class BaseModelService(ABC):
return highest_by_base
def _paginate(self, data: List[Dict], page: int, page_size: int) -> Dict:
def _paginate(self, data: List[Dict[str, Any]], page: int, page_size: int) -> Dict[str, Any]:
"""Apply pagination to filtered data"""
total_items = len(data)
start_idx = (page - 1) * page_size
@@ -791,7 +845,7 @@ class BaseModelService(ABC):
}
@abstractmethod
async def format_response(self, model_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Format model data for API response - must be implemented by subclasses.
Subclasses should return None for corrupted entries so the handler
@@ -800,17 +854,17 @@ class BaseModelService(ABC):
pass
# Common service methods that delegate to scanner
async def get_top_tags(self, limit: int = 20) -> List[Dict]:
async def get_top_tags(self, limit: int = 20) -> List[Dict[str, Any]]:
"""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]:
) -> List[Dict[str, Any]]:
"""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]:
async def get_base_models(self, limit: int = 20) -> List[Dict[str, Any]]:
"""Get base models sorted by frequency"""
return await self.scanner.get_base_models(limit)
@@ -877,7 +931,7 @@ class BaseModelService(ABC):
"""Get model root directories"""
return self.scanner.get_model_roots()
def filter_civitai_data(self, data: Dict, minimal: bool = False) -> Dict:
def filter_civitai_data(self, data: Dict[str, Any], minimal: bool = False) -> Dict[str, Any]:
"""Filter relevant fields from CivitAI data"""
if not data:
return {}
@@ -903,7 +957,7 @@ class BaseModelService(ABC):
)
return {k: data[k] for k in fields if k in data}
async def get_folder_tree(self, model_root: str) -> Dict:
async def get_folder_tree(self, model_root: str) -> Dict[str, Any]:
"""Get hierarchical folder tree for a specific model root"""
cache = await self.scanner.get_cached_data()
@@ -932,7 +986,7 @@ class BaseModelService(ABC):
return tree
async def get_unified_folder_tree(self) -> Dict:
async def get_unified_folder_tree(self) -> Dict[str, Any]:
"""Get unified folder tree across all model roots"""
cache = await self.scanner.get_cached_data()
@@ -961,7 +1015,7 @@ class BaseModelService(ABC):
return unified_tree
async def get_model_notes(self, model_name: str) -> Optional[dict]:
async def get_model_notes(self, model_name: str) -> Optional[dict[str, Any]]:
"""Get notes and file_path for a specific model file.
Supports both simple names (``OWSMianne_ANIMA_V1``) and full-path
@@ -1093,7 +1147,7 @@ class BaseModelService(ABC):
return {"civitai_url": None, "model_id": None, "version_id": None}
async def get_model_metadata(self, file_path: str) -> Optional[Dict]:
async def get_model_metadata(self, file_path: str) -> Optional[Dict[str, Any]]:
"""Load full metadata for a single model.
Listing/search endpoints return lightweight cache entries; this method performs
@@ -1189,7 +1243,7 @@ class BaseModelService(ABC):
return True
@staticmethod
def _relative_path_sort_key(relative_path: str, include_terms: List[str]) -> tuple:
def _relative_path_sort_key(relative_path: str, include_terms: List[str]) -> tuple[int, int, int, str]:
"""Sort paths by how well they satisfy the include tokens.
Sorts based on path without extension for consistent ordering.
@@ -1216,19 +1270,87 @@ class BaseModelService(ABC):
)
async def search_relative_paths(
self, search_term: str, limit: int = 15, offset: int = 0
self,
search_term: str,
limit: int = 15,
offset: int = 0,
*,
folder: Optional[str] = None,
folder_include: Optional[list[str]] = None,
folder_exclude: Optional[list[str]] = None,
base_models: Optional[list[str]] = None,
model_types: Optional[list[str]] = None,
tags: Optional[dict[str, str]] = None,
auto_tags: Optional[dict[str, str]] = None,
tag_logic: str = "any",
credit_required: Optional[bool] = None,
allow_selling_generated_content: Optional[bool] = None,
recursive: bool = True,
apply_filters: bool = False,
) -> List[str]:
"""Search model relative file paths for autocomplete functionality"""
"""Search model relative file paths for autocomplete functionality.
Optional filter kwargs mirror the filters used by the list endpoint
(/api/lm/{prefix}/list). When no filter kwargs are provided the
behavior is identical to plain token-based path matching.
"""
cache = await self.scanner.get_cached_data()
include_terms, exclude_terms = self._parse_search_tokens(search_term)
data = cache.raw_data
has_filters = any(
[
apply_filters,
folder is not None,
folder_include,
folder_exclude,
base_models,
model_types,
tags,
auto_tags,
credit_required is not None,
allow_selling_generated_content is not None,
]
)
if has_filters:
# Auto-tags are not stored in the scanner cache — they are computed
# on the fly. Pre-compute them only when an auto-tag filter is
# active to avoid mutating cache entries unnecessarily.
if auto_tags:
from .auto_tag_service import extract_auto_tags
for item in data:
if not item.get("auto_tags"):
item["auto_tags"] = extract_auto_tags(item)
criteria = FilterCriteria(
folder=folder,
folder_include=folder_include,
folder_exclude=folder_exclude,
base_models=base_models,
model_types=model_types,
tags=tags,
auto_tags=auto_tags,
search_options={"recursive": recursive},
tag_logic=tag_logic,
)
data = self.filter_set.apply(data, criteria)
if credit_required is not None:
data = await self._apply_credit_required_filter(
data, credit_required
)
if allow_selling_generated_content is not None:
data = await self._apply_allow_selling_filter(
data, allow_selling_generated_content
)
matching_paths = []
# Get model roots for path calculation
model_roots = self.scanner.get_model_roots()
# Collect all matching paths first (needed for proper sorting and offset)
for model in cache.raw_data:
for model in data:
file_path = model.get("file_path", "")
if not file_path:
continue
+30 -2
View File
@@ -59,6 +59,7 @@ class CacheEntryValidator:
'notes': ('', False),
'usage_tips': ('', False),
'hash_status': ('completed', False),
'autov3': (None, False),
}
@classmethod
@@ -119,8 +120,13 @@ class CacheEntryValidator:
if is_required:
errors.append(f"Required field '{field_name}' is missing or None")
if auto_repair:
working_entry[field_name] = cls._get_default_copy(default_value)
repaired = True
# A missing optional field whose default is None is already
# semantically equal to its default (e.g. autov3: absent
# means "not checked") — writing None back is a no-op, not
# a repair.
if default_value is not None:
working_entry[field_name] = cls._get_default_copy(default_value)
repaired = True
continue
# Validate field type and value
@@ -175,6 +181,15 @@ class CacheEntryValidator:
# that invalidates the entry, but we also don't mark it repaired.
pass
# Normalize autov3 to lowercase if needed (optional field, never stripped).
autov3 = working_entry.get('autov3')
if isinstance(autov3, str) and autov3:
normalized_autov3 = autov3.lower()
if normalized_autov3 != autov3:
if auto_repair:
working_entry['autov3'] = normalized_autov3
repaired = True
# Determine if entry is valid
# Entry is valid if no critical required field errors remain after repair
# Critical fields are file_path and sha256
@@ -242,6 +257,19 @@ class CacheEntryValidator:
"""
expected_type = type(default_value)
# Special case: autov3 is optional with a three-state contract.
# None = not checked, "" = checked but unavailable, otherwise a
# 12-character hex string (case-insensitive here; normalized to
# lowercase separately).
if field_name == 'autov3':
if value is None or value == "":
return None
if not isinstance(value, str):
return f"Field 'autov3' should be string or None, got {type(value).__name__}"
if len(value) != 12 or any(c not in '0123456789abcdefABCDEF' for c in value):
return "Field 'autov3' should be a 12-character hex string"
return None
# Special handling for numeric types
if expected_type == int:
if not isinstance(value, (int, float)):
+33 -6
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import asyncio
import json
import logging
@@ -6,7 +10,7 @@ from datetime import datetime
from typing import Any, Dict, List, Optional
from ..utils.models import CheckpointMetadata
from ..utils.file_utils import find_preview_file, normalize_path
from ..utils.file_utils import find_preview_file, normalize_path, calculate_autov3
from ..utils.metadata_manager import MetadataManager
from ..config import config
from .model_scanner import ModelScanner
@@ -62,6 +66,11 @@ class CheckpointScanner(ModelScanner):
# Find preview image
preview_url = find_preview_file(base_name, dir_path)
# AutoV3 reads only the safetensors header, so it is cheap even for
# large checkpoints; record the checked state at creation time ("" =
# checked but unavailable).
autov3 = calculate_autov3(real_path)
# Create metadata WITHOUT calculating hash
metadata = CheckpointMetadata(
file_name=base_name,
@@ -77,6 +86,7 @@ class CheckpointScanner(ModelScanner):
sub_type="checkpoint",
from_civitai=False, # Mark as local model since no hash yet
hash_status="pending", # Mark hash as pending
autov3=autov3 or "",
)
# Save the created metadata
@@ -120,7 +130,11 @@ class CheckpointScanner(ModelScanner):
# 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)
self._hash_index.add_entry(
metadata.sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
async with self._hash_calculation_lock:
@@ -132,7 +146,11 @@ class CheckpointScanner(ModelScanner):
and metadata.hash_status == "completed"
and metadata.sha256
):
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
self._hash_index.add_entry(
metadata.sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
task = self._hash_calculation_tasks.get(real_path)
@@ -185,7 +203,11 @@ class CheckpointScanner(ModelScanner):
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)
self._hash_index.add_entry(
metadata.sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
# Update status to calculating
@@ -202,7 +224,11 @@ class CheckpointScanner(ModelScanner):
await MetadataManager.save_metadata(file_path, metadata)
# Update hash index
self._hash_index.add_entry(sha256.lower(), file_path)
self._hash_index.add_entry(
sha256.lower(),
file_path,
getattr(metadata, "autov3", None) or None,
)
# Update the in-memory cache entry so that subsequent
# _persist_current_cache / _save_persistent_cache calls
@@ -216,6 +242,7 @@ class CheckpointScanner(ModelScanner):
if entry.get("file_path") == file_path:
entry["sha256"] = sha256.lower()
entry["hash_status"] = "completed"
self.bump_cache_version()
break
logger.info(f"Hash calculated for checkpoint: {file_path}")
@@ -405,7 +432,7 @@ class CheckpointScanner(ModelScanner):
roots.extend(config.extra_checkpoints_roots or [])
roots.extend(config.extra_unet_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root not in seen:
+28 -28
View File
@@ -1,6 +1,6 @@
import os
import logging
from typing import Dict, Optional
from typing import Any, Dict, Optional
from .base_model_service import BaseModelService
from .auto_tag_service import extract_auto_tags
@@ -21,58 +21,58 @@ class CheckpointService(BaseModelService):
"""
super().__init__("checkpoint", scanner, CheckpointMetadata, update_service=update_service)
async def format_response(self, checkpoint_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""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")
file_path = model_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>"),
model_data.get("file_name", "<unknown>"),
)
return None
# Get sub_type from cache entry (new canonical field)
sub_type = checkpoint_data.get("sub_type", "checkpoint")
sub_type = model_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 ""
file_name = model_data.get("file_name") or ""
model_name = model_data.get("model_name") or file_name
folder = model_data.get("folder") or ""
return {
"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", ""),
"preview_url": config.get_preview_static_url(model_data.get("preview_url", "")),
"preview_nsfw_level": model_data.get("preview_nsfw_level", 0),
"base_model": model_data.get("base_model", ""),
"folder": folder,
"sha256": checkpoint_data.get("sha256", ""),
"sha256": model_data.get("sha256", ""),
"file_path": file_path.replace(os.sep, "/"),
"file_size": checkpoint_data.get("size", 0),
"modified": checkpoint_data.get("modified", ""),
"tags": checkpoint_data.get("tags", []),
"from_civitai": checkpoint_data.get("from_civitai", True),
"usage_count": checkpoint_data.get("usage_count", 0),
"notes": checkpoint_data.get("notes", ""),
"file_size": model_data.get("size", 0),
"modified": model_data.get("modified", ""),
"tags": model_data.get("tags", []),
"from_civitai": model_data.get("from_civitai", True),
"usage_count": model_data.get("usage_count", 0),
"notes": model_data.get("notes", ""),
"sub_type": sub_type,
"favorite": checkpoint_data.get("favorite", False),
"exclude": bool(checkpoint_data.get("exclude", False)),
"update_available": bool(checkpoint_data.get("update_available", False)),
"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", ""),
"favorite": model_data.get("favorite", False),
"exclude": bool(model_data.get("exclude", False)),
"update_available": bool(model_data.get("update_available", False)),
"skip_metadata_refresh": bool(model_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(model_data.get("civitai", {}), minimal=True),
"auto_tags": model_data.get("auto_tags") or extract_auto_tags(model_data),
"version_count": model_data.get("version_count"),
"hf_url": model_data.get("hf_url", ""),
}
def find_duplicate_hashes(self) -> Dict:
def find_duplicate_hashes(self) -> Dict[str, Any]:
"""Find Checkpoints with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict:
def find_duplicate_filenames(self) -> Dict[str, Any]:
"""Find Checkpoints with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()
+41 -36
View File
@@ -1,8 +1,12 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import json
import logging
import asyncio
from copy import deepcopy
from typing import Optional, Dict, Tuple, List
from typing import Any, Optional, Dict, Tuple, List, cast
from .model_metadata_provider import CivArchiveModelMetadataProvider, ModelMetadataProviderManager
from .downloader import get_downloader
from .errors import RateLimitError
@@ -37,8 +41,8 @@ class CivArchiveClient:
async def _request_json(
self,
path: str,
params: Optional[Dict[str, str]] = None
) -> Tuple[Optional[Dict], Optional[str]]:
params: Optional[Dict[str, Any]] = None
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Call CivArchive API and return JSON payload"""
success, payload = await self._make_request(path, params=params)
if not success:
@@ -52,12 +56,12 @@ class CivArchiveClient:
self,
path: str,
*,
params: Optional[Dict[str, str]] = None,
) -> Tuple[bool, Dict | str]:
params: Optional[Dict[str, Any]] = None,
) -> Tuple[bool, Dict[str, Any] | str]:
"""Wrapper around downloader.make_request that surfaces rate limits."""
downloader = await get_downloader()
kwargs: Dict[str, Dict[str, str]] = {}
kwargs: Dict[str, Dict[str, Any]] = {}
if params:
safe_params = {str(key): str(value) for key, value in params.items() if value is not None}
if safe_params:
@@ -73,10 +77,11 @@ class CivArchiveClient:
if payload.provider is None:
payload.provider = "civarchive_api"
raise payload
return success, payload
# RateLimitError is always raised above, so the returned payload is a dict or str.
return success, cast(Dict[str, Any] | str, payload)
@staticmethod
def _normalize_payload(payload: Dict) -> Dict:
def _normalize_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
"""Unwrap CivArchive responses that wrap content under a data key"""
if not isinstance(payload, dict):
return {}
@@ -86,12 +91,12 @@ class CivArchiveClient:
return payload
@staticmethod
def _split_context(payload: Dict) -> Tuple[Dict, Dict, List[Dict]]:
def _split_context(payload: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any], List[Dict[str, Any]]]:
"""Separate version payload from surrounding model context"""
data = CivArchiveClient._normalize_payload(payload)
context: Dict = {}
fallback_files: List[Dict] = []
version: Dict = {}
context: Dict[str, Any] = {}
fallback_files: List[Dict[str, Any]] = []
version: Dict[str, Any] = {}
for key, value in data.items():
if key in {"version", "model"}:
@@ -115,7 +120,7 @@ class CivArchiveClient:
return context, version, fallback_files
@staticmethod
def _ensure_list(value) -> List:
def _ensure_list(value: Any) -> List[Any]:
if isinstance(value, list):
return value
if value is None:
@@ -123,7 +128,7 @@ class CivArchiveClient:
return [value]
@staticmethod
def _build_model_info(context: Dict) -> Dict:
def _build_model_info(context: Dict[str, Any]) -> Dict[str, Any]:
tags = context.get("tags")
if not isinstance(tags, list):
tags = list(tags) if isinstance(tags, (set, tuple)) else ([] if tags is None else [tags])
@@ -136,7 +141,7 @@ class CivArchiveClient:
}
@staticmethod
def _build_creator_info(context: Dict) -> Dict:
def _build_creator_info(context: Dict[str, Any]) -> Dict[str, Any]:
username = context.get("creator_username") or context.get("username") or ""
image = context.get("creator_image") or context.get("creator_avatar") or ""
creator: Dict[str, Optional[str]] = {
@@ -150,7 +155,7 @@ class CivArchiveClient:
return creator
@staticmethod
def _transform_file_entry(file_data: Dict) -> Dict:
def _transform_file_entry(file_data: Dict[str, Any]) -> Dict[str, Any]:
mirrors = file_data.get("mirrors") or []
if not isinstance(mirrors, list):
mirrors = [mirrors]
@@ -165,7 +170,7 @@ class CivArchiveClient:
if not name and available_mirror:
name = available_mirror.get("filename")
transformed: Dict = {
transformed: Dict[str, Any] = {
"id": file_data.get("id"),
"sizeKB": file_data.get("sizeKB"),
"name": name,
@@ -216,23 +221,23 @@ class CivArchiveClient:
def _transform_files(
self,
files: Optional[List[Dict]],
fallback_files: Optional[List[Dict]] = None
) -> List[Dict]:
candidates: List[Dict] = []
files: Optional[List[Dict[str, Any]]],
fallback_files: Optional[List[Dict[str, Any]]] = None
) -> List[Dict[str, Any]]:
candidates: List[Dict[str, Any]] = []
if isinstance(files, list) and files:
candidates = files
elif isinstance(fallback_files, list):
candidates = fallback_files
transformed_files: List[Dict] = []
transformed_files: List[Dict[str, Any]] = []
for file_data in candidates:
if isinstance(file_data, dict):
transformed_files.append(self._transform_file_entry(file_data))
# Sort: .safetensors first, .ckpt second, others last
# so the backend fallback (no file_params) prefers safetensors
def _sort_key(f: Dict) -> int:
def _sort_key(f: Dict[str, Any]) -> int:
fname = f.get("name") or ""
if isinstance(fname, str):
lower = fname.lower()
@@ -247,10 +252,10 @@ class CivArchiveClient:
def _transform_version(
self,
context: Dict,
version: Dict,
fallback_files: Optional[List[Dict]] = None
) -> Optional[Dict]:
context: Dict[str, Any],
version: Dict[str, Any],
fallback_files: Optional[List[Dict[str, Any]]] = None
) -> Optional[Dict[str, Any]]:
if not version:
return None
@@ -291,7 +296,7 @@ class CivArchiveClient:
return version_copy
async def _resolve_version_from_files(self, payload: Dict) -> Optional[Dict]:
async def _resolve_version_from_files(self, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Fallback to fetch version data when only file metadata is available"""
data = self._normalize_payload(payload)
files = data.get("files") or payload.get("files") or []
@@ -323,7 +328,7 @@ class CivArchiveClient:
return resolved
return None
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by SHA256 hash value using CivArchive API"""
try:
payload, error = await self._request_json(f"/sha256/{model_hash.lower()}")
@@ -332,12 +337,12 @@ class CivArchiveClient:
return None, "Model not found"
return None, error
context, version_data, fallback_files = self._split_context(payload)
context, version_data, fallback_files = self._split_context(cast(Dict[str, Any], payload))
transformed = self._transform_version(context, version_data, fallback_files)
if transformed:
return transformed, None
resolved = await self._resolve_version_from_files(payload)
resolved = await self._resolve_version_from_files(cast(Dict[str, Any], payload))
if resolved:
return resolved, None
@@ -350,7 +355,7 @@ class CivArchiveClient:
logger.error(f"Error fetching CivArchive model by hash {model_hash[:10]}: {e}")
return None, str(e)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model using CivArchive API"""
try:
payload, error = await self._request_json(f"/models/{model_id}")
@@ -364,7 +369,7 @@ class CivArchiveClient:
context, version_data, fallback_files = self._split_context(payload)
versions_meta = data.get("versions") or []
transformed_versions: List[Dict] = []
transformed_versions: List[Dict[str, Any]] = []
for meta in versions_meta:
if not isinstance(meta, dict):
continue
@@ -381,7 +386,7 @@ class CivArchiveClient:
if primary_version:
transformed_versions.insert(0, primary_version)
ordered_versions: List[Dict] = []
ordered_versions: List[Dict[str, Any]] = []
seen_ids = set()
for version in transformed_versions:
version_id = version.get("id")
@@ -402,7 +407,7 @@ class CivArchiveClient:
logger.error(f"Error fetching CivArchive model versions for {model_id}: {e}")
return None
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: int | str | None = None, version_id: int | str | None = None) -> Optional[Dict[str, Any]]:
"""Get specific model version using CivArchive API
Args:
@@ -459,7 +464,7 @@ class CivArchiveClient:
logger.error(f"Error fetching CivArchive model version via API {model_id}/{version_id}: {e}")
return None
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
""" Fetch model version metadata using a known bogus model lookup
CivArchive lacks a direct version lookup API, this uses a workaround (which we handle in the main model request now)
+1 -1
View File
@@ -283,7 +283,7 @@ class CivitaiBaseModelService:
return None
if isinstance(result, str):
data = json.loads(result)
data: Any = json.loads(result)
else:
data = result
+128 -40
View File
@@ -1,9 +1,14 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import asyncio
import copy
import logging
import os
import time
from collections import OrderedDict
from typing import Any, Optional, Dict, Tuple, List, Sequence
from typing import Any, Optional, Dict, Tuple, List, Sequence, cast
from .connectivity_guard import (
OFFLINE_FRIENDLY_MESSAGE,
is_expected_offline_error,
@@ -19,6 +24,12 @@ from ..utils.civitai_utils import resolve_license_payload
logger = logging.getLogger(__name__)
# Best-effort cache for creator model counts, keyed by lowercase username.
# Values are (monotonic timestamp, count or None); None results are cached
# too so repeated failures don't hammer the API.
_CREATOR_COUNT_CACHE_TTL_SECONDS = 600
_creator_model_count_cache: Dict[str, Tuple[float, Optional[int]]] = {}
class CivitaiClient:
_instance = None
@@ -51,7 +62,7 @@ class CivitaiClient:
# Uses OrderedDict with LRU eviction at MAX_CACHE_ENTRIES to prevent
# unbounded growth in long-running server processes.
self._version_info_cache: OrderedDict[
str, Tuple[Optional[Dict], Optional[str]]
str, Tuple[Optional[Dict[str, Any]], Optional[str]]
] = OrderedDict()
self._MAX_CACHE_ENTRIES = 500
@@ -65,7 +76,7 @@ class CivitaiClient:
*,
use_auth: bool = False,
**kwargs,
) -> Tuple[bool, Dict | str]:
) -> Tuple[bool, Dict[str, Any] | str]:
"""Wrapper around downloader.make_request that surfaces rate limits,
with retry for transient server errors (5xx, Cloudflare 524, network flakiness)."""
@@ -79,7 +90,8 @@ class CivitaiClient:
**kwargs,
)
if success:
return True, result
# RateLimitError is raised below; a successful result is dict or str.
return True, cast(Dict[str, Any] | str, result)
if isinstance(result, RateLimitError):
if result.provider is None:
@@ -119,7 +131,7 @@ class CivitaiClient:
return False, "Unexpected error in _make_request"
@staticmethod
def _remove_comfy_metadata(model_version: Optional[Dict]) -> None:
def _remove_comfy_metadata(model_version: Optional[Dict[str, Any]]) -> None:
"""Remove Comfy-specific metadata from model version images."""
if not isinstance(model_version, dict):
return
@@ -166,7 +178,7 @@ class CivitaiClient:
async def get_model_by_hash(
self, model_hash: str
) -> Tuple[Optional[Dict], Optional[str]]:
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
try:
success, version = await self._make_request(
"GET",
@@ -213,7 +225,7 @@ class CivitaiClient:
# Ensure directory exists
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
f.write(content)
f.write(content if isinstance(content, bytes) else content.encode("utf-8"))
return True
return False
except Exception as e:
@@ -268,7 +280,7 @@ class CivitaiClient:
return True
return False
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model with local availability info"""
try:
success, result = await self._make_request(
@@ -276,7 +288,7 @@ class CivitaiClient:
f"{self.base_url}/models/{model_id}",
use_auth=True,
)
if success:
if success and isinstance(result, dict):
# Also return model type along with versions
return {
"modelVersions": result.get("modelVersions", []),
@@ -310,7 +322,7 @@ class CivitaiClient:
async def get_model_versions_bulk(
self, model_ids: Sequence[int]
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
"""Fetch model metadata for multiple ids using the batch API."""
deduped: Dict[int, None] = {}
@@ -340,13 +352,13 @@ class CivitaiClient:
if not isinstance(items, list):
return {}
payload: Dict[int, Dict] = {}
payload: Dict[int, Dict[str, Any]] = {}
for item in items:
if not isinstance(item, dict):
continue
model_id = item.get("id")
try:
normalized_id = int(model_id)
normalized_id = int(cast(Any, model_id))
except (TypeError, ValueError):
continue
payload[normalized_id] = {
@@ -366,8 +378,8 @@ class CivitaiClient:
return None
async def get_model_version(
self, model_id: int = None, version_id: int = None
) -> Optional[Dict]:
self, model_id: int | None = None, version_id: int | None = None
) -> Optional[Dict[str, Any]]:
"""Get specific model version with additional metadata."""
try:
if model_id is None and version_id is not None:
@@ -385,7 +397,7 @@ class CivitaiClient:
logger.error(f"Error fetching model version: {e}")
return None
async def _get_version_by_id_only(self, version_id: int) -> Optional[Dict]:
async def _get_version_by_id_only(self, version_id: int) -> Optional[Dict[str, Any]]:
version = await self._fetch_version_by_id(version_id)
if version is None:
return None
@@ -404,7 +416,7 @@ class CivitaiClient:
async def _get_version_with_model_id(
self, model_id: int, version_id: Optional[int]
) -> Optional[Dict]:
) -> Optional[Dict[str, Any]]:
model_data = await self._fetch_model_data(model_id)
if not model_data:
return None
@@ -457,20 +469,20 @@ class CivitaiClient:
self._remove_comfy_metadata(version)
return version
async def _fetch_model_data(self, model_id: int) -> Optional[Dict]:
async def _fetch_model_data(self, model_id: int) -> Optional[Dict[str, Any]]:
success, data = await self._make_request(
"GET",
f"{self.base_url}/models/{model_id}",
use_auth=True,
)
if success:
if success and isinstance(data, dict):
return data
if is_expected_offline_error(data):
return None
logger.warning(f"Failed to fetch model data for model {model_id}")
return None
async def _fetch_version_by_id(self, version_id: Optional[int]) -> Optional[Dict]:
async def _fetch_version_by_id(self, version_id: Optional[int]) -> Optional[Dict[str, Any]]:
if version_id is None:
return None
@@ -479,7 +491,7 @@ class CivitaiClient:
f"{self.base_url}/model-versions/{version_id}",
use_auth=True,
)
if success:
if success and isinstance(version, dict):
return version
if is_expected_offline_error(version):
return None
@@ -487,7 +499,7 @@ class CivitaiClient:
logger.warning(f"Failed to fetch version by id {version_id}")
return None
async def _fetch_version_by_hash(self, model_hash: Optional[str]) -> Optional[Dict]:
async def _fetch_version_by_hash(self, model_hash: Optional[str]) -> Optional[Dict[str, Any]]:
if not model_hash:
return None
@@ -496,7 +508,7 @@ class CivitaiClient:
f"{self.base_url}/model-versions/by-hash/{model_hash}",
use_auth=True,
)
if success:
if success and isinstance(version, dict):
return version
if is_expected_offline_error(version):
return None
@@ -505,8 +517,8 @@ class CivitaiClient:
return None
def _select_target_version(
self, model_data: Dict, model_id: int, version_id: Optional[int]
) -> Optional[Dict]:
self, model_data: Dict[str, Any], model_id: int, version_id: Optional[int]
) -> Optional[Dict[str, Any]]:
model_versions = model_data.get("modelVersions", [])
if not model_versions:
logger.warning(f"No model versions found for model {model_id}")
@@ -525,7 +537,7 @@ class CivitaiClient:
return model_versions[0]
def _extract_primary_model_hash(self, version_entry: Dict) -> Optional[str]:
def _extract_primary_model_hash(self, version_entry: Dict[str, Any]) -> Optional[str]:
for file_info in version_entry.get("files", []):
if file_info.get("type") == "Model" and file_info.get("primary"):
hashes = file_info.get("hashes", {})
@@ -535,8 +547,8 @@ class CivitaiClient:
return None
def _build_version_from_model_data(
self, version_entry: Dict, model_id: int, model_data: Dict
) -> Dict:
self, version_entry: Dict[str, Any], model_id: int, model_data: Dict[str, Any]
) -> Dict[str, Any]:
version = copy.deepcopy(version_entry)
version.pop("index", None)
version["modelId"] = model_id
@@ -548,7 +560,7 @@ class CivitaiClient:
}
return version
def _enrich_version_with_model_data(self, version: Dict, model_data: Dict) -> None:
def _enrich_version_with_model_data(self, version: Dict[str, Any], model_data: Dict[str, Any]) -> None:
model_info = version.get("model")
if not isinstance(model_info, dict):
model_info = {}
@@ -564,7 +576,7 @@ class CivitaiClient:
async def get_model_version_info(
self, version_id: str
) -> Tuple[Optional[Dict], Optional[str]]:
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version metadata from Civitai
Args:
@@ -589,7 +601,7 @@ class CivitaiClient:
logger.debug("Resolving Civitai model version info: %s", url)
success, result = await self._make_request("GET", url, use_auth=True)
if success:
if success and isinstance(result, dict):
logger.debug("Successfully fetched model version info for: %s", version_id)
self._remove_comfy_metadata(result)
self._version_info_cache[version_id] = (result, None)
@@ -619,7 +631,7 @@ class CivitaiClient:
async def get_image_info(
self, image_id: str, source_url: str | None = None
) -> Optional[Dict]:
) -> Optional[Dict[str, Any]]:
"""Fetch image information from Civitai API
Args:
@@ -652,7 +664,7 @@ class CivitaiClient:
)
return None
if result and "items" in result and isinstance(result["items"], list):
if isinstance(result, dict) and "items" in result and isinstance(result["items"], list):
items = result["items"]
for item in items:
@@ -692,7 +704,7 @@ class CivitaiClient:
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
"""Fetch full version details for up to 100 SHA256 hashes via the batch endpoint.
Uses POST /api/v1/model-versions/by-hash which returns full version
@@ -709,7 +721,7 @@ class CivitaiClient:
return []
BATCH_SIZE = 100
all_versions: List[Dict] = []
all_versions: List[Dict[str, Any]] = []
for start in range(0, len(hashes), BATCH_SIZE):
batch = hashes[start : start + BATCH_SIZE]
@@ -729,7 +741,7 @@ class CivitaiClient:
continue
if isinstance(result, list):
all_versions.extend(result)
all_versions.extend(cast(Any, result))
else:
logger.debug(
"Unexpected by-hash response type: %s", type(result)
@@ -743,17 +755,34 @@ class CivitaiClient:
return all_versions if all_versions else None
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
"""Fetch all models for a specific Civitai user."""
async def get_user_models(
self, username: str, cursor: Optional[str] = None
) -> Optional[Dict[str, Any]]:
"""Fetch one page (up to 100 models) for a specific Civitai user.
Returns ``{"items": [...], "nextCursor": <str|None>}`` on success,
or None on failure. Pass ``cursor`` (from a previous response's
``nextCursor``) to fetch subsequent pages.
"""
if not username:
return None
params: Dict[str, Any] = {
"username": username,
"nsfw": "true",
"limit": 100,
"sort": "Newest",
"period": "AllTime",
}
if cursor:
params["cursor"] = cursor
try:
success, result = await self._make_request(
"GET",
f"{self.base_url}/models",
use_auth=True,
params={"username": username, "nsfw": "true"},
params=params,
)
if not success:
@@ -765,7 +794,7 @@ class CivitaiClient:
items = result.get("items") if isinstance(result, dict) else None
if not isinstance(items, list):
return []
items = []
for model in items:
versions = model.get("modelVersions")
@@ -774,9 +803,68 @@ class CivitaiClient:
for version in versions:
self._remove_comfy_metadata(version)
return items
next_cursor: Optional[str] = None
metadata = result.get("metadata") if isinstance(result, dict) else None
if isinstance(metadata, dict):
raw_cursor = metadata.get("nextCursor")
if raw_cursor is not None:
next_cursor = str(raw_cursor)
return {"items": items, "nextCursor": next_cursor}
except RateLimitError:
raise
except Exception as exc: # pragma: no cover - defensive logging
logger.error("Error fetching models for %s: %s", username, exc)
return None
async def get_creator_model_count(self, username: str) -> Optional[int]:
"""Best-effort lookup of a creator's published model count.
Uses the ``/creators`` endpoint (a contains-match query), picking the
entry whose username matches exactly (case-insensitive). Returns None
on any failure; never raises. Results (including None) are cached
for ``_CREATOR_COUNT_CACHE_TTL_SECONDS``.
"""
if not username:
return None
cache_key = username.lower()
cached = _creator_model_count_cache.get(cache_key)
if cached is not None:
cached_at, cached_count = cached
if time.monotonic() - cached_at < _CREATOR_COUNT_CACHE_TTL_SECONDS:
return cached_count
count: Optional[int] = None
try:
success, result = await self._make_request(
"GET",
f"{self.base_url}/creators",
use_auth=True,
params={"query": username, "limit": 10},
)
if success and isinstance(result, dict):
creators = result.get("items")
if isinstance(creators, list):
for creator in creators:
if not isinstance(creator, dict):
continue
creator_name = creator.get("username")
if not isinstance(creator_name, str):
continue
if creator_name.lower() != cache_key:
continue
model_count = creator.get("modelCount")
if isinstance(model_count, (int, float)) and not isinstance(
model_count, bool
):
count = int(model_count)
break
except Exception as exc: # best-effort only, never propagate
logger.debug(
"Failed to fetch creator model count for %s: %s", username, exc
)
_creator_model_count_cache[cache_key] = (time.monotonic(), count)
return count
+1 -1
View File
@@ -18,7 +18,7 @@ class DownloadCoordinator:
self,
*,
ws_manager,
download_manager_factory: Callable[[], Awaitable],
download_manager_factory: Callable[[], Awaitable[Any]],
) -> None:
self._ws_manager = ws_manager
self._download_manager_factory = download_manager_factory
+135 -76
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import copy
import logging
import os
@@ -8,7 +12,7 @@ import zipfile
from concurrent.futures import ThreadPoolExecutor
from collections import OrderedDict
import uuid
from typing import Dict, List, Optional, Set, Tuple
from typing import Any, Dict, List, Optional, Set, Tuple, cast
from urllib.parse import urlparse
from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
from ..utils.constants import (
@@ -18,7 +22,7 @@ from ..utils.constants import (
VALID_LORA_TYPES,
)
from ..utils.civitai_utils import normalize_civitai_download_url, rewrite_preview_url
from ..utils.file_utils import calculate_sha256
from ..utils.file_utils import calculate_sha256, calculate_autov3
from ..utils.preview_selection import resolve_mature_threshold, select_preview_media
from ..utils.utils import sanitize_folder_name
from ..utils.exif_utils import ExifUtils
@@ -121,7 +125,7 @@ class DownloadManager:
"delay": 0,
}
)
except DownloadInProgressError:
except DownloadInProgressError: # pyright: ignore[reportPossiblyUnboundVariable]
logger.info(
"Skipping automatic example images download for %s; another example images download is already running",
model_hash,
@@ -170,7 +174,7 @@ class DownloadManager:
logger.error("aria2 download failed for %s: %s", download_url, exc)
return False, str(exc)
download_kwargs = {
download_kwargs: Dict[str, Any] = {
"progress_callback": progress_callback,
"use_auth": use_auth,
}
@@ -204,16 +208,16 @@ class DownloadManager:
async def download_from_civitai(
self,
model_id: int = None,
model_version_id: int = None,
save_dir: str = None,
model_id: int | None = None,
model_version_id: int | None = None,
save_dir: str | None = None,
relative_path: str = "",
progress_callback=None,
use_default_paths: bool = False,
download_id: str = None,
source: str = None,
file_params: Dict = None,
) -> Dict:
download_id: str | None = None,
source: str | None = None,
file_params: Dict[str, Any] | None = None,
) -> Dict[str, Any]:
"""Download model from Civitai with task tracking and concurrency control
Args:
@@ -309,14 +313,14 @@ class DownloadManager:
async def _download_with_semaphore(
self,
task_id: str,
model_id: int,
model_version_id: int,
save_dir: str,
model_id: int | None,
model_version_id: int | None,
save_dir: str | None,
relative_path: str,
progress_callback=None,
use_default_paths: bool = False,
source: str = None,
file_params: Dict = None,
source: str | None = None,
file_params: Dict[str, Any] | None = None,
):
"""Execute download with semaphore to limit concurrency"""
# Update status to waiting
@@ -380,7 +384,8 @@ class DownloadManager:
# Use original download implementation
try:
# Check for cancellation before starting
if asyncio.current_task().cancelled():
current_task = asyncio.current_task()
if current_task is not None and current_task.cancelled():
raise asyncio.CancelledError()
result = await self._execute_original_download(
@@ -484,11 +489,11 @@ class DownloadManager:
# Schedule cleanup of download record after delay
asyncio.create_task(self._cleanup_download_record(task_id))
def _start_background_download_task(self, download_id: str, coroutine) -> asyncio.Task:
def _start_background_download_task(self, download_id: str, coroutine) -> asyncio.Task[Any]:
task = asyncio.create_task(coroutine)
self._download_tasks[download_id] = task
def _cleanup_done_task(done_task: asyncio.Task) -> None:
def _cleanup_done_task(done_task: asyncio.Task[Any]) -> None:
current_task = self._download_tasks.get(download_id)
if current_task is done_task:
self._download_tasks.pop(download_id, None)
@@ -530,7 +535,7 @@ class DownloadManager:
async def _cleanup_cancelled_download_files(
self,
download_id: str,
download_info: Optional[Dict],
download_info: Optional[Dict[str, Any]],
) -> None:
target_files = set()
persisted = await self._aria2_state_store.get(download_id)
@@ -603,13 +608,13 @@ class DownloadManager:
self,
download_id: str,
*,
extra: Optional[Dict] = None,
extra: Optional[Dict[str, Any]] = None,
) -> None:
info = self._active_downloads.get(download_id)
if not info:
return
payload = {
payload: Dict[str, Any] = {
"download_id": download_id,
"model_id": info.get("model_id"),
"model_version_id": info.get("model_version_id"),
@@ -631,7 +636,7 @@ class DownloadManager:
await self._aria2_state_store.upsert(download_id, payload)
def _build_restored_download_info(self, record: Dict, save_path: str) -> Dict:
def _build_restored_download_info(self, record: Dict[str, Any], save_path: str) -> Dict[str, Any]:
return {
"model_id": record.get("model_id"),
"model_version_id": record.get("model_version_id"),
@@ -653,8 +658,8 @@ class DownloadManager:
def _is_same_aria2_download_request(
self,
current_info: Optional[Dict],
persisted_record: Dict,
current_info: Optional[Dict[str, Any]],
persisted_record: Dict[str, Any],
) -> bool:
if not isinstance(current_info, dict):
return False
@@ -666,13 +671,15 @@ class DownloadManager:
return current_version_id == persisted_version_id
def _build_download_urls_from_file_info(self, file_info: Dict, source: str = None) -> List[str]:
def _build_download_urls_from_file_info(self, file_info: Dict[str, Any], source: str | None = None) -> List[str]:
mirrors = file_info.get("mirrors") or []
download_urls: List[str] = []
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"]))
normalized_url = normalize_civitai_download_url(mirror["url"])
if normalized_url:
download_urls.append(normalized_url)
if source == "civarchive" and len(download_urls) > 1:
civitai_urls = [
@@ -688,7 +695,9 @@ class DownloadManager:
if not download_urls:
download_url = file_info.get("downloadUrl")
if download_url:
download_urls.append(normalize_civitai_download_url(download_url))
normalized_url = normalize_civitai_download_url(download_url)
if normalized_url:
download_urls.append(normalized_url)
return download_urls
@@ -696,8 +705,8 @@ class DownloadManager:
self,
*,
model_type: str,
version_info: Dict,
file_info: Dict,
version_info: Dict[str, Any],
file_info: Dict[str, Any],
save_path: str,
):
if model_type == "checkpoint":
@@ -706,7 +715,7 @@ class DownloadManager:
return EmbeddingMetadata.from_civitai_info(version_info, file_info, save_path)
return LoraMetadata.from_civitai_info(version_info, file_info, save_path)
def _resolve_save_path_from_persisted_record(self, record: Dict) -> Optional[str]:
def _resolve_save_path_from_persisted_record(self, record: Dict[str, Any]) -> Optional[str]:
save_path = record.get("save_path") or record.get("file_path")
if isinstance(save_path, str) and save_path:
return os.path.abspath(save_path)
@@ -728,7 +737,7 @@ class DownloadManager:
return os.path.abspath(os.path.join(save_dir, file_name))
async def _resume_restored_aria2_download(self, download_id: str, record: Dict) -> Dict:
async def _resume_restored_aria2_download(self, download_id: str, record: Dict[str, Any]) -> Dict[str, Any]:
try:
if download_id in self._active_downloads:
self._active_downloads[download_id]["status"] = "downloading"
@@ -842,7 +851,7 @@ class DownloadManager:
self,
previous_download_id: str,
new_download_id: str,
persisted_record: Dict,
persisted_record: Dict[str, Any],
save_path: str,
) -> None:
aria2_downloader = await get_aria2_downloader()
@@ -938,7 +947,7 @@ class DownloadManager:
except Exception:
status_payload = None
if status_payload is not None:
if status_payload is not None and isinstance(gid, str):
remote_status = status_payload.get("status", "")
if remote_status in {"active", "waiting", "paused"}:
await aria2_downloader.restore_transfer(download_id, gid, save_path)
@@ -1115,17 +1124,17 @@ class DownloadManager:
async def _execute_original_download(
self,
model_id,
model_version_id,
save_dir,
relative_path,
model_id: int | None,
model_version_id: int | None,
save_dir: str | None,
relative_path: str,
progress_callback,
use_default_paths,
download_id=None,
transfer_backend="python",
source=None,
file_params=None,
):
use_default_paths: bool,
download_id: str | None = None,
transfer_backend: str = "python",
source: str | None = None,
file_params: Dict[str, Any] | None = None,
) -> Dict[str, Any]:
"""Wrapper for original download_from_civitai implementation"""
try:
# Check if model version already exists in library
@@ -1172,7 +1181,7 @@ class DownloadManager:
# Get version info based on the provided identifier
version_info = await metadata_provider.get_model_version(
model_id, model_version_id
cast(int, model_id), cast(int, model_version_id)
)
if not version_info:
@@ -1183,7 +1192,7 @@ class DownloadManager:
)
metadata_provider = await get_default_metadata_provider()
version_info = await metadata_provider.get_model_version(
model_id, model_version_id
cast(int, model_id), cast(int, model_version_id)
)
if not version_info:
@@ -1388,8 +1397,20 @@ class DownloadManager:
relative_path = self._calculate_relative_path(version_info, model_type)
# Update save directory with relative path if provided
if not save_dir:
return {"success": False, "error": "No save directory specified"}
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)
@@ -1551,6 +1572,11 @@ class DownloadManager:
version_info, file_info, save_path
)
logger.info(f"Creating EmbeddingMetadata for {file_name}")
else:
return {
"success": False,
"error": f'Unsupported model type "{model_type}"',
}
# 6. Start download process
if transfer_backend == "aria2" and download_id:
@@ -1570,7 +1596,7 @@ class DownloadManager:
},
)
execute_kwargs = {
execute_kwargs: Dict[str, Any] = {
"download_urls": download_urls,
"save_dir": save_dir,
"metadata": metadata,
@@ -1617,7 +1643,8 @@ class DownloadManager:
)
# If early_access_msg exists and download failed, replace error message
if "early_access_msg" in locals() and not result.get("success", False):
early_access_msg = locals().get("early_access_msg")
if early_access_msg and not result.get("success", False):
result["error"] = early_access_msg
return result
@@ -1642,7 +1669,7 @@ class DownloadManager:
self,
model_type: str,
model_id_value,
version_info: Dict,
version_info: Dict[str, Any],
fallback_version_id=None,
file_path: str | None = None,
) -> None:
@@ -1673,8 +1700,8 @@ class DownloadManager:
try:
await history_service.mark_downloaded(
model_type,
int(version_id),
model_id=int(resolved_model_id) if resolved_model_id is not None else None,
int(cast(Any, version_id)),
model_id=int(cast(Any, resolved_model_id)) if resolved_model_id is not None else None,
source="download",
file_path=file_path,
)
@@ -1691,7 +1718,7 @@ class DownloadManager:
self,
model_type: str,
model_id_value,
version_info: Dict,
version_info: Dict[str, Any],
fallback_version_id=None,
) -> None:
"""Ensure update tracking reflects a newly downloaded version."""
@@ -1715,7 +1742,7 @@ class DownloadManager:
if isinstance(model_info, dict):
resolved_model_id = model_info.get("id")
try:
resolved_model_id = int(resolved_model_id)
resolved_model_id = int(cast(Any, resolved_model_id))
except (TypeError, ValueError):
logger.debug(
"Skipping update sync; invalid model id: %s", resolved_model_id
@@ -1726,7 +1753,7 @@ class DownloadManager:
if version_id is None:
version_id = fallback_version_id
try:
version_id = int(version_id)
version_id = int(cast(Any, version_id))
except (TypeError, ValueError):
logger.debug(
"Skipping update sync; invalid version id for model %s: %s",
@@ -1763,7 +1790,7 @@ class DownloadManager:
for entry in local_versions or []:
vid = entry.get("versionId")
try:
version_ids.add(int(vid))
version_ids.add(int(cast(Any, vid)))
except (TypeError, ValueError):
continue
@@ -1785,7 +1812,7 @@ class DownloadManager:
)
def _calculate_relative_path(
self, version_info: Dict, model_type: str = "lora"
self, version_info: Dict[str, Any], model_type: str = "lora"
) -> str:
"""Calculate relative path using template from settings
@@ -1827,6 +1854,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)
@@ -1842,6 +1872,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(
@@ -1849,21 +1888,22 @@ class DownloadManager:
download_urls: List[str],
save_dir: str,
metadata,
version_info: Dict,
version_info: Dict[str, Any],
relative_path: str,
progress_callback=None,
model_type: str = "lora",
download_id: str = None,
download_id: str | None = None,
transfer_backend: Optional[str] = None,
) -> Dict:
) -> Dict[str, Any]:
"""Execute the actual download process including preview images and model files"""
metadata_entries: List = []
metadata_entries: List[Any] = []
metadata_files_for_cleanup: List[str] = []
extracted_paths: List[str] = []
metadata_path = ""
preview_targets: List[str] = []
preview_path: str | None = None
preview_nsfw_level = 0
save_path: str | None = None
transfer_backend = (transfer_backend or self._get_model_download_backend()).lower()
try:
resolved, save_path = await self._resolve_download_target_path(
@@ -1911,9 +1951,9 @@ class DownloadManager:
mature_threshold=mature_threshold,
)
preview_url = selected_image.get("url") if selected_image else None
preview_url = cast(Optional[str], selected_image.get("url")) if selected_image else None
media_type = (
(selected_image.get("type") or "").lower() if selected_image else ""
cast(str, selected_image.get("type") or "").lower() if selected_image else ""
)
def _extension_from_url(url: str, fallback: str) -> str:
@@ -1937,9 +1977,10 @@ class DownloadManager:
preview_url, media_type="video"
)
attempt_urls: List[str] = []
if rewritten:
if rewritten and rewritten_url:
attempt_urls.append(rewritten_url)
attempt_urls.append(preview_url)
if preview_url:
attempt_urls.append(preview_url)
seen_attempts = set()
for attempt in attempt_urls:
@@ -1956,7 +1997,7 @@ class DownloadManager:
rewritten_url, rewritten = rewrite_preview_url(
preview_url, media_type="image"
)
if rewritten:
if rewritten and rewritten_url:
preview_ext = _extension_from_url(preview_url, ".png")
preview_path = os.path.splitext(save_path)[0] + preview_ext
success, _ = await downloader.download_file(
@@ -1982,7 +2023,9 @@ class DownloadManager:
)
if success:
with open(temp_path, "wb") as temp_file_handle:
temp_file_handle.write(content)
temp_file_handle.write(
content if isinstance(content, bytes) else content.encode("utf-8")
)
preview_path = (
os.path.splitext(save_path)[0] + ".webp"
)
@@ -2034,6 +2077,8 @@ class DownloadManager:
last_error = None
for download_url in download_urls:
download_url = normalize_civitai_download_url(download_url)
if download_url is None:
continue
use_auth = download_url.startswith(CIVITAI_DOWNLOAD_URL_PREFIXES)
if transfer_backend == "aria2" and download_id:
await self._persist_aria2_state(
@@ -2138,6 +2183,10 @@ class DownloadManager:
"error": f"Zip archive does not contain any supported model files ({supported_text})",
}
actual_file_paths = extracted_paths
# The archive entry's AutoV3 (if any) describes the zip itself,
# not the extracted models; clear it so per-file header
# resolution applies to every extracted model.
metadata.autov3 = None
try:
os.remove(save_path)
except OSError as exc:
@@ -2213,7 +2262,7 @@ class DownloadManager:
entry, normalized_file_path, adjust_root
)
if adjusted_entry is not None:
entry = adjusted_entry
entry = cast(Any, adjusted_entry)
metadata_entries[index] = entry
metadata_file_path = (
@@ -2333,11 +2382,11 @@ class DownloadManager:
async def _build_metadata_entries(
self, base_metadata, file_paths: List[str]
) -> List:
) -> List[Any]:
if not file_paths:
return []
entries: List = []
entries: List[Any] = []
for index, file_path in enumerate(file_paths):
entry = base_metadata if index == 0 else copy.deepcopy(base_metadata)
# Update file paths without modifying size and modified timestamps
@@ -2352,6 +2401,16 @@ class DownloadManager:
sha256 = await calculate_sha256(file_path)
if sha256:
entry.sha256 = sha256.lower()
# AutoV3: the Civitai-reported value for the downloaded file (set
# by from_civitai_info) takes precedence. Only the un-checked
# state (None) triggers a header read; '' (checked-unavailable)
# is never re-read, honoring the three-state contract so rows
# marked at download time stay untouched by later passes.
if entry.autov3 is None:
autov3 = await asyncio.get_running_loop().run_in_executor(
None, calculate_autov3, file_path
)
entry.autov3 = (autov3 or "").lower()
entries.append(entry)
return entries
@@ -2370,7 +2429,7 @@ class DownloadManager:
return destination
def _distribute_preview_to_entries(
self, preview_path: str, entries: List
self, preview_path: str, entries: List[Any]
) -> List[str]:
if not preview_path or not entries:
return []
@@ -2429,7 +2488,7 @@ class DownloadManager:
progress_callback, normalized_snapshot, rounded_progress
)
async def cancel_download(self, download_id: str) -> Dict:
async def cancel_download(self, download_id: str) -> Dict[str, Any]:
"""Cancel an active download by download_id
Args:
@@ -2511,7 +2570,7 @@ class DownloadManager:
self._download_tasks.pop(download_id, None)
await self._aria2_state_store.remove(download_id)
async def skip_download(self, download_id: str) -> Dict:
async def skip_download(self, download_id: str) -> Dict[str, Any]:
"""Skip a download while preserving all partial files on disk.
Removes all in-memory tracking (asyncio task, semaphore, active/pause
@@ -2594,7 +2653,7 @@ class DownloadManager:
# Preserve aria2 state store entry so the partial download
# info survives restarts and can be resumed later
async def pause_download(self, download_id: str) -> Dict:
async def pause_download(self, download_id: str) -> Dict[str, Any]:
"""Pause an active download without losing progress."""
await self._restore_persisted_downloads()
@@ -2641,7 +2700,7 @@ class DownloadManager:
return {"success": True, "message": "Download paused successfully"}
async def resume_download(self, download_id: str) -> Dict:
async def resume_download(self, download_id: str) -> Dict[str, Any]:
"""Resume a previously paused download."""
await self._restore_persisted_downloads()
@@ -2658,7 +2717,7 @@ class DownloadManager:
self._pause_events[download_id] = pause_control
self._active_downloads[download_id] = self._build_restored_download_info(
persisted,
os.path.abspath(save_path),
os.path.abspath(cast(str, save_path)),
)
if pause_control.is_set():
@@ -2785,7 +2844,7 @@ class DownloadManager:
elif asyncio.iscoroutine(result):
await result
async def get_active_downloads(self) -> Dict:
async def get_active_downloads(self) -> Dict[str, Any]:
"""Get information about all active downloads
Returns:
+34 -2
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,
@@ -74,6 +74,9 @@ class DownloadQueueService:
);
CREATE INDEX IF NOT EXISTS idx_dh_completed ON download_history(completed_at DESC);
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;
"""
@@ -115,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
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
from __future__ import annotations
import asyncio
+13 -8
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""
Unified download manager for all HTTP/HTTPS downloads in the application.
@@ -20,7 +24,7 @@ from dataclasses import dataclass
from datetime import datetime, timedelta
from email.utils import parsedate_to_datetime
from urllib.parse import urlparse
from typing import Optional, Dict, Tuple, Callable, Union, Awaitable
from typing import Optional, Dict, Tuple, Callable, Union, Awaitable, Any, cast
from ..services.settings_manager import get_settings_manager
from .connectivity_guard import (
OFFLINE_COOLDOWN_ERROR,
@@ -204,6 +208,7 @@ class Downloader:
# Double check after acquiring lock
if self._session is None or self._should_refresh_session():
await self._create_session()
assert self._session is not None
return self._session
@property
@@ -231,7 +236,7 @@ class Downloader:
)
try:
timeout_value = float(raw_value)
timeout_value = float(cast(Any, raw_value))
except (TypeError, ValueError):
timeout_value = default_timeout
@@ -243,7 +248,7 @@ class Downloader:
raw_value = os.environ.get("COMFYUI_DOWNLOAD_MAX_RETRIES")
try:
retries = int(raw_value)
retries = int(cast(Any, raw_value))
except (TypeError, ValueError):
retries = default_retries
@@ -320,7 +325,7 @@ class Downloader:
# CA coverage across different Python environments (especially
# embedded/compatibility Python builds).
try:
import certifi # type: ignore[import-untyped]
import certifi # pyright: ignore[reportMissingTypeStubs]
ca_path = certifi.where()
ssl_context = ssl.create_default_context(cafile=ca_path)
@@ -330,7 +335,7 @@ class Downloader:
logger.debug("SSL: certifi unavailable; using system default CA bundle")
# Optimize TCP connection parameters
connector_kwargs = dict(
connector_kwargs: Dict[str, Any] = dict(
ssl=ssl_context,
limit=8, # Concurrent connections
ttl_dns_cache=300, # DNS cache timeout
@@ -890,7 +895,7 @@ class Downloader:
use_auth: bool = False,
custom_headers: Optional[Dict[str, str]] = None,
return_headers: bool = False,
) -> Tuple[bool, Union[bytes, str], Optional[Dict]]:
) -> Tuple[bool, Union[bytes, str], Optional[Dict[str, Any]]]:
"""
Download a file to memory (for small files like preview images)
@@ -976,7 +981,7 @@ class Downloader:
url: str,
use_auth: bool = False,
custom_headers: Optional[Dict[str, str]] = None,
) -> Tuple[bool, Union[Dict, str]]:
) -> Tuple[bool, Union[Dict[str, Any], str]]:
"""
Get response headers without downloading the full content
@@ -1036,7 +1041,7 @@ class Downloader:
use_auth: bool = False,
custom_headers: Optional[Dict[str, str]] = None,
**kwargs,
) -> Tuple[bool, Union[Dict, str]]:
) -> Tuple[bool, Union[Dict[str, Any], str, RateLimitError]]:
"""
Make a generic HTTP request and return JSON response
+1 -1
View File
@@ -27,7 +27,7 @@ class EmbeddingScanner(ModelScanner):
roots.extend(config.embeddings_roots or [])
roots.extend(config.extra_embeddings_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root and root not in seen:
+28 -28
View File
@@ -1,6 +1,6 @@
import os
import logging
from typing import Dict, Optional
from typing import Any, Dict, Optional
from .base_model_service import BaseModelService
from .auto_tag_service import extract_auto_tags
@@ -21,58 +21,58 @@ class EmbeddingService(BaseModelService):
"""
super().__init__("embedding", scanner, EmbeddingMetadata, update_service=update_service)
async def format_response(self, embedding_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""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")
file_path = model_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>"),
model_data.get("file_name", "<unknown>"),
)
return None
# Get sub_type from cache entry (new canonical field)
sub_type = embedding_data.get("sub_type", "embedding")
sub_type = model_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 ""
file_name = model_data.get("file_name") or ""
model_name = model_data.get("model_name") or file_name
folder = model_data.get("folder") or ""
return {
"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", ""),
"preview_url": config.get_preview_static_url(model_data.get("preview_url", "")),
"preview_nsfw_level": model_data.get("preview_nsfw_level", 0),
"base_model": model_data.get("base_model", ""),
"folder": folder,
"sha256": embedding_data.get("sha256", ""),
"sha256": model_data.get("sha256", ""),
"file_path": file_path.replace(os.sep, "/"),
"file_size": embedding_data.get("size", 0),
"modified": embedding_data.get("modified", ""),
"tags": embedding_data.get("tags", []),
"from_civitai": embedding_data.get("from_civitai", True),
# "usage_count": embedding_data.get("usage_count", 0), # TODO: Enable when embedding usage tracking is implemented
"notes": embedding_data.get("notes", ""),
"file_size": model_data.get("size", 0),
"modified": model_data.get("modified", ""),
"tags": model_data.get("tags", []),
"from_civitai": model_data.get("from_civitai", True),
# "usage_count": model_data.get("usage_count", 0), # TODO: Enable when embedding usage tracking is implemented
"notes": model_data.get("notes", ""),
"sub_type": sub_type,
"favorite": embedding_data.get("favorite", False),
"exclude": bool(embedding_data.get("exclude", False)),
"update_available": bool(embedding_data.get("update_available", False)),
"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", ""),
"favorite": model_data.get("favorite", False),
"exclude": bool(model_data.get("exclude", False)),
"update_available": bool(model_data.get("update_available", False)),
"skip_metadata_refresh": bool(model_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(model_data.get("civitai", {}), minimal=True),
"auto_tags": model_data.get("auto_tags") or extract_auto_tags(model_data),
"version_count": model_data.get("version_count"),
"hf_url": model_data.get("hf_url", ""),
}
def find_duplicate_hashes(self) -> Dict:
def find_duplicate_hashes(self) -> Dict[str, Any]:
"""Find Embeddings with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict:
def find_duplicate_filenames(self) -> Dict[str, Any]:
"""Find Embeddings with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()
@@ -35,7 +35,7 @@ class CleanupResult:
def to_dict(self) -> Dict[str, object]:
"""Convert the dataclass to a serialisable dictionary."""
data = {
data: Dict[str, object] = {
"success": self.success,
"checked_folders": self.checked_folders,
"moved_empty_folders": self.moved_empty_folders,
+47 -8
View File
@@ -201,6 +201,11 @@ PROVIDER_PRESETS: Dict[str, Dict[str, Any]] = {
"api_base": "https://openrouter.ai/api/v1",
"requires_key": True,
},
"google": {
"name": "Gemini",
"api_base": "https://generativelanguage.googleapis.com/v1beta/openai",
"requires_key": True,
},
"opencode-go": {
"name": "OpenCode Go",
"api_base": "https://opencode.ai/zen/go/v1",
@@ -566,18 +571,52 @@ class LLMService:
if effective_max is None:
effective_max = 4096
result = await self.chat_completion(
messages=messages,
model=model,
temperature=temperature,
response_format={"type": "json_object"},
max_tokens=effective_max,
)
# 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 in json_object mode. "
"LLM returned empty content. "
f"Raw response: {json.dumps(result)[:500]}"
)
+14 -4
View File
@@ -1,10 +1,12 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import logging
from typing import List
from ..utils.models import LoraMetadata
from ..config import config
from .model_scanner import ModelScanner
from .model_hash_index import ModelHashIndex # Changed from LoraHashIndex to ModelHashIndex
import sys
logger = logging.getLogger(__name__)
@@ -15,8 +17,10 @@ class LoraScanner(ModelScanner):
def __init__(self):
# Define supported file extensions
file_extensions = {'.safetensors'}
# Initialize parent class with ModelHashIndex
from .model_hash_index import ModelHashIndex
super().__init__(
model_type="lora",
model_class=LoraMetadata,
@@ -26,11 +30,13 @@ class LoraScanner(ModelScanner):
def get_model_roots(self) -> List[str]:
"""Get lora root directories (including extra paths)"""
from ..config import config
roots: List[str] = []
roots.extend(config.loras_roots or [])
roots.extend(config.extra_loras_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root and root not in seen:
@@ -68,8 +74,12 @@ class LoraScanner(ModelScanner):
test_hash = next(iter(self._hash_index._hash_to_path.keys()))
test_path = self._hash_index.get_path(test_hash)
logger.debug(f"\nTest lookup by hash: {test_hash[:8]}... -> {test_path}")
if test_path is None:
return
# Also test reverse lookup
test_hash_result = self._hash_index.get_hash(test_path)
if test_hash_result is None:
return
logger.debug(f"Test reverse lookup: {test_path} -> {test_hash_result[:8]}...\n\n")
+41 -41
View File
@@ -1,7 +1,7 @@
import logging
import json
import os
from typing import Dict, List, Optional
from typing import Any, Dict, List, Optional
from .base_model_service import BaseModelService
from .model_query import resolve_sub_type
@@ -24,7 +24,7 @@ class LoraService(BaseModelService):
"""
super().__init__("lora", scanner, LoraMetadata, update_service=update_service)
async def format_response(self, lora_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Format LoRA data for API response.
Returns None when the entry is missing critical fields (corrupted cache
@@ -32,56 +32,56 @@ class LoraService(BaseModelService):
whole listing request. See issue #730.
"""
# Guard against corrupted cache entries missing critical fields
file_path = lora_data.get("file_path")
file_path = model_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>"),
model_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()
sub_type = resolve_sub_type(model_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 ""
file_name = model_data.get("file_name") or ""
model_name = model_data.get("model_name") or file_name
folder = model_data.get("folder") or ""
return {
"model_name": model_name,
"file_name": file_name,
"preview_url": config.get_preview_static_url(
lora_data.get("preview_url", "")
model_data.get("preview_url", "")
),
"preview_nsfw_level": lora_data.get("preview_nsfw_level", 0),
"base_model": lora_data.get("base_model", ""),
"preview_nsfw_level": model_data.get("preview_nsfw_level", 0),
"base_model": model_data.get("base_model", ""),
"folder": folder,
"sha256": lora_data.get("sha256", ""),
"sha256": model_data.get("sha256", ""),
"file_path": file_path.replace(os.sep, "/"),
"file_size": lora_data.get("size", 0),
"modified": lora_data.get("modified", ""),
"tags": lora_data.get("tags", []),
"from_civitai": lora_data.get("from_civitai", True),
"usage_count": lora_data.get("usage_count", 0),
"usage_tips": lora_data.get("usage_tips", ""),
"notes": lora_data.get("notes", ""),
"favorite": lora_data.get("favorite", False),
"exclude": bool(lora_data.get("exclude", False)),
"update_available": bool(lora_data.get("update_available", False)),
"file_size": model_data.get("size", 0),
"modified": model_data.get("modified", ""),
"tags": model_data.get("tags", []),
"from_civitai": model_data.get("from_civitai", True),
"usage_count": model_data.get("usage_count", 0),
"usage_tips": model_data.get("usage_tips", ""),
"notes": model_data.get("notes", ""),
"favorite": model_data.get("favorite", False),
"exclude": bool(model_data.get("exclude", False)),
"update_available": bool(model_data.get("update_available", False)),
"skip_metadata_refresh": bool(
lora_data.get("skip_metadata_refresh", False)
model_data.get("skip_metadata_refresh", False)
),
"sub_type": sub_type,
"civitai": self.filter_civitai_data(
lora_data.get("civitai", {}), minimal=True
model_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", ""),
"auto_tags": model_data.get("auto_tags") or extract_auto_tags(model_data),
"version_count": model_data.get("version_count"),
"hf_url": model_data.get("hf_url", ""),
}
async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]:
async def _apply_specific_filters(self, data: List[Dict[str, Any]], **kwargs) -> List[Dict[str, Any]]:
"""Apply LoRA-specific filters"""
# Handle first_letter filter for LoRAs
first_letter = kwargs.get("first_letter")
@@ -152,7 +152,7 @@ class LoraService(BaseModelService):
return data
def _filter_by_first_letter(self, data: List[Dict], letter: str) -> List[Dict]:
def _filter_by_first_letter(self, data: List[Dict[str, Any]], letter: str) -> List[Dict[str, Any]]:
"""Filter data by first letter of model name
Special handling:
@@ -307,7 +307,7 @@ class LoraService(BaseModelService):
return None
@staticmethod
def get_recommended_strength_from_lora_data(lora_data: Dict) -> Optional[float]:
def get_recommended_strength_from_lora_data(lora_data: Dict[str, Any]) -> Optional[float]:
"""Parse usage_tips JSON and extract recommended model strength."""
try:
usage_tips = lora_data.get("usage_tips", "")
@@ -320,7 +320,7 @@ class LoraService(BaseModelService):
@staticmethod
def get_recommended_clip_strength_from_lora_data(
lora_data: Dict,
lora_data: Dict[str, Any],
) -> Optional[float]:
"""Parse usage_tips JSON and extract recommended clip strength."""
try:
@@ -332,7 +332,7 @@ class LoraService(BaseModelService):
except (json.JSONDecodeError, TypeError, AttributeError):
return None
async def get_lora_metadata_by_filename(self, filename: str) -> Optional[Dict]:
async def get_lora_metadata_by_filename(self, filename: str) -> Optional[Dict[str, Any]]:
"""Return cached raw metadata for a LoRA matching the given filename."""
cache = await self.scanner.get_cached_data(force_refresh=False)
@@ -357,11 +357,11 @@ class LoraService(BaseModelService):
return None
def find_duplicate_hashes(self) -> Dict:
def find_duplicate_hashes(self) -> Dict[str, Any]:
"""Find LoRAs with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict:
def find_duplicate_filenames(self) -> Dict[str, Any]:
"""Find LoRAs with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()
@@ -373,8 +373,8 @@ class LoraService(BaseModelService):
use_same_clip_strength: bool = True,
clip_strength_min: float = 0.0,
clip_strength_max: float = 1.0,
locked_loras: Optional[List[Dict]] = None,
pool_config: Optional[Dict] = None,
locked_loras: Optional[List[Dict[str, Any]]] = None,
pool_config: Optional[Dict[str, Any]] = None,
count_mode: str = "fixed",
count_min: int = 3,
count_max: int = 7,
@@ -382,7 +382,7 @@ class LoraService(BaseModelService):
recommended_strength_scale_min: float = 0.5,
recommended_strength_scale_max: float = 1.0,
seed: Optional[int] = None,
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""
Get random LoRAs with specified strength ranges.
@@ -513,8 +513,8 @@ class LoraService(BaseModelService):
return result_loras
async def _apply_pool_filters(
self, available_loras: List[Dict], pool_config: Dict
) -> List[Dict]:
self, available_loras: List[Dict[str, Any]], pool_config: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""
Apply pool_config filters to available LoRAs.
@@ -671,8 +671,8 @@ class LoraService(BaseModelService):
return available_loras
async def get_cycler_list(
self, pool_config: Optional[Dict] = None, sort_by: str = "filename"
) -> List[Dict]:
self, pool_config: Optional[Dict[str, Any]] = None, sort_by: str = "filename"
) -> List[Dict[str, Any]]:
"""
Get filtered and sorted LoRA list for cycling.
+5 -1
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import os
import logging
from .model_metadata_provider import (
@@ -170,7 +174,7 @@ def _wrap_provider_with_rate_limit(provider_name: str | None, provider: ModelMet
return RateLimitRetryingProvider(provider, label=provider_name)
async def get_metadata_provider(provider_name: str = None):
async def get_metadata_provider(provider_name: str | None = None):
"""Get a specific metadata provider or default provider with rate-limit handling."""
provider_manager = await ModelMetadataProviderManager.get_instance()
+20 -7
View File
@@ -6,25 +6,26 @@ import json
import logging
import os
from datetime import datetime
from typing import Any, Awaitable, Callable, Dict, Iterable, Optional
from typing import Any, Awaitable, Callable, Dict, Iterable, Optional, Protocol
from ..services.settings_manager import SettingsManager
from ..utils.civitai_utils import resolve_license_payload
from ..utils.model_utils import determine_base_model
from ..utils.models import autov3_from_civitai_files
from .connectivity_guard import OFFLINE_FRIENDLY_MESSAGE, is_expected_offline_error
from .errors import RateLimitError
logger = logging.getLogger(__name__)
class MetadataProviderProtocol:
class MetadataProviderProtocol(Protocol):
"""Subset of metadata provider interface consumed by the sync service."""
async def get_model_by_hash(self, sha256: str) -> tuple[Optional[Dict[str, Any]], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> tuple[Optional[Dict[str, Any]], Optional[str]]:
...
async def get_model_version(
self, model_id: int, model_version_id: Optional[int]
self, model_id: Any = None, version_id: Any = None
) -> Optional[Dict[str, Any]]:
...
@@ -38,8 +39,8 @@ class MetadataSyncService:
metadata_manager,
preview_service,
settings: SettingsManager,
default_metadata_provider_factory: Callable[[], Awaitable[MetadataProviderProtocol]],
metadata_provider_selector: Callable[[str], Awaitable[MetadataProviderProtocol]],
default_metadata_provider_factory: Callable[..., Awaitable[MetadataProviderProtocol]],
metadata_provider_selector: Callable[..., Awaitable[MetadataProviderProtocol]],
) -> None:
self._metadata_manager = metadata_manager
self._preview_service = preview_service
@@ -152,6 +153,18 @@ class MetadataSyncService:
civitai_metadata.get("baseModel")
)
# Civitai-first AutoV3 propagation: the freshly fetched version
# metadata may report an AutoV3 for the file whose SHA256 matches the
# local model. Persist it now so recipe matching sees it immediately —
# no full rescan or restart required (the header is never re-read to
# upgrade the checked-unavailable '' state).
sha256_value = (local_metadata.get("sha256") or "").lower()
civitai_autov3 = autov3_from_civitai_files(
local_metadata.get("civitai"), sha256_value
)
if civitai_autov3:
local_metadata["autov3"] = civitai_autov3
await self._preview_service.ensure_preview_for_metadata(
metadata_path, local_metadata, civitai_metadata.get("images", [])
)
@@ -479,7 +492,7 @@ class MetadataSyncService:
if not file_paths:
raise ValueError("No file paths provided for verification")
results = {
results: Dict[str, Any] = {
"verified_as_duplicates": True,
"mismatched_files": [],
"new_hash_map": {},
+35 -21
View File
@@ -1,6 +1,7 @@
import asyncio
import time
import logging
import random
logger = logging.getLogger(__name__)
from typing import Any, Dict, List, Optional, Tuple
@@ -30,17 +31,22 @@ DISPLAY_NAME_MODES = {"model_name", "file_name"}
class ModelCache:
"""Cache structure for model data with extensible sorting."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
folders: List[str]
version_index: Dict[int, Dict] = field(default_factory=dict)
version_index: Dict[int, Dict[str, Any]] = field(default_factory=dict)
model_id_index: Dict[int, List[Dict[str, Any]]] = field(default_factory=dict)
name_display_mode: str = "model_name"
_lock: Any = field(init=False, repr=False, default=None)
# Cache for last sort: (sort_key, order, seed) -> sorted list
_last_sort: Tuple[Optional[str], str, Optional[str]] = field(
init=False, repr=False, default=(None, "asc", None)
)
_last_sorted_data: List[Dict[str, Any]] = field(
init=False, repr=False, default_factory=list
)
def __post_init__(self):
self._lock = asyncio.Lock()
# Cache for last sort: (sort_key, order) -> sorted list
self._last_sort: Tuple[str, str] = (None, None)
self._last_sorted_data: List[Dict] = []
self._normalize_raw_data()
self.name_display_mode = self._normalize_display_mode(self.name_display_mode)
# Default sort on init
@@ -63,7 +69,7 @@ class ModelCache:
return ""
return str(value)
def _normalize_item(self, item: Dict) -> None:
def _normalize_item(self, item: Dict[str, Any]) -> None:
"""Ensure core metadata fields are present and string typed."""
if not isinstance(item, dict):
@@ -79,7 +85,7 @@ class ModelCache:
for item in self.raw_data:
self._normalize_item(item)
def _get_display_name(self, item: Dict) -> str:
def _get_display_name(self, item: Dict[str, Any]) -> str:
"""Return the value used for name-based sorting based on display settings."""
if self.name_display_mode == "file_name":
@@ -113,7 +119,7 @@ class ModelCache:
for item in self.raw_data:
self.add_to_version_index(item)
def add_to_version_index(self, item: Dict) -> None:
def add_to_version_index(self, item: Dict[str, Any]) -> None:
"""Register a cache item in the version/model indexes if possible."""
civitai_data = item.get('civitai') if isinstance(item, dict) else None
@@ -142,7 +148,7 @@ class ModelCache:
else:
versions.append(descriptor)
def remove_from_version_index(self, item: Dict) -> None:
def remove_from_version_index(self, item: Dict[str, Any]) -> None:
"""Remove a cache item from the version/model indexes if present."""
civitai_data = item.get('civitai') if isinstance(item, dict) else None
@@ -176,7 +182,7 @@ class ModelCache:
def _build_version_descriptor(
self,
item: Dict,
item: Dict[str, Any],
civitai_data: Dict[str, Any],
version_id: int,
) -> Optional[Dict[str, Any]]:
@@ -203,9 +209,9 @@ class ModelCache:
async def resort(self):
"""Resort cached data according to last sort mode if set"""
async with self._lock:
if self._last_sort != (None, None):
sort_key, order = self._last_sort
sorted_data = self._sort_data(self.raw_data, sort_key, order)
sort_key, order, seed = self._last_sort
if sort_key is not None:
sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
self._last_sorted_data = sorted_data
# Update folder list
# else: do nothing
@@ -218,7 +224,7 @@ class ModelCache:
self.folders = sorted(list(all_folders), key=lambda x: x.lower())
self.rebuild_version_index()
def _sort_data(self, data: List[Dict], sort_key: str, order: str) -> List[Dict]:
def _sort_data(self, data: List[Dict[str, Any]], sort_key: str, order: str, seed: Optional[str] = None) -> List[Dict[str, Any]]:
"""Sort data by sort_key and order"""
start_time = time.perf_counter()
reverse = (order == 'desc')
@@ -265,6 +271,13 @@ class ModelCache:
),
reverse=reverse
)
elif sort_key == 'random':
# Random shuffle seeded for stable pagination: the same seed
# always yields the same order, so successive page requests
# stay consistent while browsing.
rng = random.Random(seed or 'random')
result = list(data)
rng.shuffle(result)
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.
@@ -285,15 +298,16 @@ class ModelCache:
logger.debug("ModelCache._sort_data(%s, %s) for %d items took %.3fs", sort_key, order, len(data), duration)
return result
async def get_sorted_data(self, sort_key: str = 'name', order: str = 'asc') -> List[Dict]:
async def get_sorted_data(self, sort_key: str = 'name', order: str = 'asc', seed: Optional[str] = None) -> List[Dict[str, Any]]:
"""Get sorted data by sort_key and order, using cache if possible"""
async with self._lock:
if (sort_key, order) == self._last_sort:
cache_key = (sort_key, order, seed)
if cache_key == self._last_sort:
return self._last_sorted_data
start_time = time.perf_counter()
sorted_data = self._sort_data(self.raw_data, sort_key, order)
self._last_sort = (sort_key, order)
sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
self._last_sort = cache_key
self._last_sorted_data = sorted_data
duration = time.perf_counter() - start_time
@@ -312,9 +326,9 @@ class ModelCache:
self.name_display_mode = normalized
if self._last_sort[0] == 'name':
sort_key, order = self._last_sort
self._last_sorted_data = self._sort_data(self.raw_data, sort_key, order)
sort_key, order, seed = self._last_sort
if sort_key == 'name':
self._last_sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
async def update_preview_url(self, file_path: str, preview_url: str, preview_nsfw_level: int) -> bool:
"""Update preview_url for a specific model in all cached data
+7 -1
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__)
@@ -40,7 +41,7 @@ class AutoOrganizeResult:
def to_dict(self) -> Dict[str, Any]:
"""Convert result to dictionary"""
result = {
result: Dict[str, Any] = {
'success': self.status != 'error',
'status': self.status,
'message': f'Auto-organize {self.operation_type} completed: {self.success_count} moved, {self.skipped_count} skipped, {self.failure_count} failed out of {self.total} total',
@@ -417,6 +418,8 @@ class ModelFileService:
"""Calculate the target directory for a model"""
if is_flat_structure:
file_path = model.get('file_path')
if not isinstance(file_path, str):
return None
current_dir = os.path.dirname(file_path)
# Check if already in root directory
@@ -493,6 +496,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()
+53 -4
View File
@@ -8,11 +8,12 @@ class ModelHashIndex:
self._hash_to_path: Dict[str, str] = {}
self._filename_to_hash: Dict[str, str] = {}
self._autov2_to_path: Dict[str, str] = {}
self._autov3_to_path: Dict[str, str] = {}
# New data structures for tracking duplicates
self._duplicate_hashes: Dict[str, List[str]] = {} # sha256 -> list of paths
self._duplicate_filenames: Dict[str, List[str]] = {} # filename -> list of paths
def add_entry(self, sha256: str, file_path: str) -> None:
def add_entry(self, sha256: str, file_path: str, autov3: Optional[str] = None) -> None:
"""Add or update hash index entry"""
if not sha256 or not file_path:
return
@@ -33,9 +34,14 @@ class ModelHashIndex:
self._duplicate_hashes.setdefault(sha256, []).append(file_path)
# Track duplicates by filename - FIXED LOGIC
is_re_registration = False
existing_hash: Optional[str] = None
if filename in self._filename_to_hash:
existing_hash = self._filename_to_hash[filename]
existing_path = self._hash_to_path.get(existing_hash)
# Same path registered again (e.g. a file replaced in place with
# new content) — used below to drop its stale autov3 mapping.
is_re_registration = existing_path == file_path
# If this is a different file with the same filename
if existing_path and existing_path != file_path:
@@ -67,12 +73,36 @@ class ModelHashIndex:
# AutoV2 = first 10 chars of SHA256
if len(sha256) >= 10:
self._autov2_to_path[sha256[:10]] = file_path
# AutoV3 is an independent hash (not derived from SHA256), stored as-is.
# Drop stale mappings for a path when it is re-registered with a NEW
# sha256 (file replaced in place) or with an explicit new autov3 value
# (correction). Re-registering the SAME file with the same sha256 and
# no autov3 (e.g. lazy-hash completion) must never clear its existing
# mapping. First-time registrations stay O(1).
if autov3:
autov3 = autov3.lower()
if is_re_registration and (existing_hash != sha256 or autov3):
stale_autov3_keys = [
key for key, mapped_path in self._autov3_to_path.items()
if mapped_path == file_path and key != autov3
]
for key in stale_autov3_keys:
del self._autov3_to_path[key]
if autov3:
self._autov3_to_path[autov3] = file_path
def add_autov3(self, autov3: str, file_path: str) -> None:
"""Add or update an AutoV3-only index entry (used when only AutoV3 is known)"""
if not autov3:
return
autov3 = autov3.lower()
self._autov3_to_path[autov3] = file_path
def _get_filename_from_path(self, file_path: str) -> str:
"""Extract filename without extension from path"""
return os.path.splitext(os.path.basename(file_path))[0]
def remove_by_path(self, file_path: str, hash_val: str = None) -> None:
def remove_by_path(self, file_path: str, hash_val: Optional[str] = None) -> None:
"""Remove entry by file path"""
filename = self._get_filename_from_path(file_path)
@@ -167,6 +197,11 @@ class ModelHashIndex:
for k in autov2_keys_to_remove:
del self._autov2_to_path[k]
# Remove from AutoV3 index
autov3_keys_to_remove = [k for k, v in self._autov3_to_path.items() if v == file_path]
for k in autov3_keys_to_remove:
del self._autov3_to_path[k]
def remove_by_hash(self, sha256: str) -> None:
"""Remove entry by hash"""
sha256 = sha256.lower()
@@ -189,6 +224,11 @@ class ModelHashIndex:
autov2_key = sha256[:10]
if autov2_key in self._autov2_to_path:
del self._autov2_to_path[autov2_key]
# Remove AutoV3 entries pointing to any removed path
autov3_keys_to_remove = [k for k, v in self._autov3_to_path.items() if v in paths_to_remove]
for k in autov3_keys_to_remove:
del self._autov3_to_path[k]
# Update filename-to-hash and duplicate filenames for all paths
for path_to_remove in paths_to_remove:
@@ -209,22 +249,26 @@ class ModelHashIndex:
del self._duplicate_filenames[fname]
def has_hash(self, hash_value: str) -> bool:
"""Check if hash exists in index (SHA256 or AutoV2)"""
"""Check if hash exists in index (SHA256, AutoV2, or AutoV3)"""
normalized = hash_value.lower()
if normalized in self._hash_to_path:
return True
if len(normalized) == 10:
return normalized in self._autov2_to_path
if len(normalized) == 12:
return normalized in self._autov3_to_path
return False
def get_path(self, hash_value: str) -> Optional[str]:
"""Get file path for a hash (SHA256 or AutoV2)"""
"""Get file path for a hash (SHA256, AutoV2, or AutoV3)"""
normalized = hash_value.lower()
path = self._hash_to_path.get(normalized)
if path is not None:
return path
if len(normalized) == 10:
return self._autov2_to_path.get(normalized)
if len(normalized) == 12:
return self._autov3_to_path.get(normalized)
return None
def get_hash(self, file_path: str) -> Optional[str]:
@@ -243,6 +287,7 @@ class ModelHashIndex:
self._hash_to_path.clear()
self._filename_to_hash.clear()
self._autov2_to_path.clear()
self._autov3_to_path.clear()
self._duplicate_hashes.clear()
self._duplicate_filenames.clear()
@@ -253,6 +298,10 @@ class ModelHashIndex:
def get_all_filenames(self) -> Set[str]:
"""Get all filenames in the index"""
return set(self._filename_to_hash.keys())
def get_all_autov3(self) -> Dict[str, str]:
"""Get a snapshot of all AutoV3 hashes mapped to their file paths"""
return dict(self._autov3_to_path)
def get_duplicate_hashes(self) -> Dict[str, List[str]]:
"""Get dictionary of duplicate hashes and their paths"""
+54 -6
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
import logging
import os
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, cast
from ..services.service_registry import ServiceRegistry
from ..utils.constants import PREVIEW_EXTENSIONS
@@ -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."""
@@ -57,8 +87,8 @@ class ModelLifecycleService:
scanner,
metadata_manager,
metadata_loader: Callable[[str], Awaitable[Dict[str, object]]],
recipe_scanner_factory: Callable[[], Awaitable] | None = None,
update_service: "ModelUpdateService" | None = None,
recipe_scanner_factory: Callable[[], Awaitable[Any]] | None = None,
update_service: Optional["ModelUpdateService"] = None,
) -> None:
self._scanner = scanner
self._metadata_manager = metadata_manager
@@ -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
@@ -106,6 +138,9 @@ class ModelLifecycleService:
item for item in cache.raw_data if item.get("file_path") != file_path
]
await cache.resort()
bump_cache_version = getattr(self._scanner, "bump_cache_version", None)
if callable(bump_cache_version):
bump_cache_version()
if hasattr(self._scanner, "_hash_index") and self._scanner._hash_index:
self._scanner._hash_index.remove_by_path(file_path)
@@ -114,7 +149,7 @@ class ModelLifecycleService:
persist_current_cache = getattr(self._scanner, "_persist_current_cache", None)
if callable(persist_current_cache):
await persist_current_cache()
await cast(Awaitable[Any], persist_current_cache())
return {"success": True, "deleted_files": deleted_files}
@@ -182,6 +217,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
@@ -210,6 +247,9 @@ class ModelLifecycleService:
item for item in cache.raw_data if item["file_path"] != file_path
]
await cache.resort()
bump_cache_version = getattr(self._scanner, "bump_cache_version", None)
if callable(bump_cache_version):
bump_cache_version()
excluded = getattr(self._scanner, "_excluded_models", None)
if isinstance(excluded, list):
@@ -218,7 +258,7 @@ class ModelLifecycleService:
persist_current_cache = getattr(self._scanner, "_persist_current_cache", None)
if callable(persist_current_cache):
await persist_current_cache()
await cast(Awaitable[Any], persist_current_cache())
message = f"Model {os.path.basename(file_path)} excluded"
return {"success": True, "message": message}
@@ -229,6 +269,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 +312,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 +325,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")
@@ -316,7 +363,8 @@ class ModelLifecycleService:
if os.path.exists(metadata_path):
metadata = await self._metadata_loader(metadata_path)
hash_value = metadata.get("sha256") if isinstance(metadata, dict) else None
raw_hash = metadata.get("sha256") if isinstance(metadata, dict) else None
hash_value = raw_hash if isinstance(raw_hash, str) else None
renamed_files: List[str] = []
new_metadata_path: Optional[str] = None
+93 -54
View File
@@ -10,7 +10,7 @@ from .errors import RateLimitError, ResourceNotFoundError
try:
from bs4 import BeautifulSoup
except ImportError as exc:
BeautifulSoup = None # type: ignore[assignment]
BeautifulSoup = None # pyright: ignore[reportAssignmentType]
_BS4_IMPORT_ERROR = exc
else:
_BS4_IMPORT_ERROR = None
@@ -18,7 +18,7 @@ else:
try:
import aiosqlite
except ImportError as exc:
aiosqlite = None # type: ignore[assignment]
aiosqlite = None # pyright: ignore[reportAssignmentType]
_AIOSQLITE_IMPORT_ERROR = exc
else:
_AIOSQLITE_IMPORT_ERROR = None
@@ -105,24 +105,24 @@ class ModelMetadataProvider(ABC):
"""Base abstract class for all model metadata providers"""
@abstractmethod
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by hash value"""
pass
@abstractmethod
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model with their details"""
pass
async def get_model_versions_bulk(
self, model_ids: Sequence[int]
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
"""Fetch model versions for multiple model ids when supported."""
raise NotImplementedError
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
"""Fetch full version details for multiple SHA256 hashes.
Used specifically to retrieve ``usageControl`` which is only
@@ -133,50 +133,61 @@ class ModelMetadataProvider(ABC):
raise NotImplementedError
@abstractmethod
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
"""Get specific model version with additional metadata"""
pass
@abstractmethod
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version metadata"""
pass
@abstractmethod
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
"""Fetch models owned by the specified user"""
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Fetch one page of models owned by the specified user.
Returns ``{"items": [...], "nextCursor": <str|None>}`` on success,
or None when unsupported/failed. ``cursor`` continues a previous page.
"""
pass
async def get_creator_model_count(self, username: str) -> Optional[int]:
"""Published model count for the user; None when unsupported."""
return None
class CivitaiModelMetadataProvider(ModelMetadataProvider):
"""Provider that uses Civitai API for metadata"""
def __init__(self, civitai_client):
self.client = civitai_client
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_by_hash(model_hash)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
return await self.client.get_model_versions(model_id)
async def get_model_versions_bulk(
self, model_ids: Sequence[int]
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
return await self.client.get_model_versions_bulk(model_ids)
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
return await self.client.get_model_versions_by_hashes(hashes)
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
return await self.client.get_model_version(model_id, version_id)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_version_info(version_id)
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
return await self.client.get_user_models(username)
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
return await self.client.get_user_models(username, cursor)
async def get_creator_model_count(self, username: str) -> Optional[int]:
return await self.client.get_creator_model_count(username)
class CivArchiveModelMetadataProvider(ModelMetadataProvider):
"""Provider that uses CivArchive API for metadata"""
@@ -184,19 +195,19 @@ class CivArchiveModelMetadataProvider(ModelMetadataProvider):
def __init__(self, civarchive_client):
self.client = civarchive_client
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_by_hash(model_hash)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
return await self.client.get_model_versions(model_id)
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
return await self.client.get_model_version(model_id, version_id)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_version_info(version_id)
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Not supported by CivArchive provider"""
return None
@@ -207,7 +218,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
self.db_path = db_path
self._aiosqlite = _require_aiosqlite()
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by hash value from SQLite database"""
async with self._aiosqlite.connect(self.db_path) as db:
# Look up in model_files table to get model_id and version_id
@@ -232,7 +243,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
result = await self._get_version_with_model_data(db, model_id, version_id)
return result, None if result else "Error retrieving model data"
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model from SQLite database"""
async with self._aiosqlite.connect(self.db_path) as db:
db.row_factory = self._aiosqlite.Row
@@ -288,7 +299,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
'name': model_name
}
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
"""Get specific model version with additional metadata from SQLite database"""
if not model_id and not version_id:
return None
@@ -328,7 +339,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
# Now we have both model_id and version_id, get the full data
return await self._get_version_with_model_data(db, model_id, version_id)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version metadata from SQLite database"""
async with self._aiosqlite.connect(self.db_path) as db:
db.row_factory = self._aiosqlite.Row
@@ -347,11 +358,11 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
version_data = await self._get_version_with_model_data(db, model_id, version_id)
return version_data, None
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Listing models by username is not supported for archive database"""
return None
async def _get_version_with_model_data(self, db, model_id, version_id) -> Optional[Dict]:
async def _get_version_with_model_data(self, db, model_id, version_id) -> Optional[Dict[str, Any]]:
"""Helper to build version data with model information"""
# Get version details
version_query = "SELECT name, base_model, data FROM model_versions WHERE id = ? AND model_id = ?"
@@ -474,7 +485,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
jitter_ratio=self._rate_limit_jitter_ratio,
)
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
for provider, label in self._iter_providers():
try:
result, error = await self._call_with_rate_limit(
@@ -496,7 +507,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None, "Model not found"
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
not_found_confirmed = False
for provider, label in self._iter_providers():
try:
@@ -527,7 +538,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
for provider, label in self._iter_providers():
try:
result = await self._call_with_rate_limit(
@@ -550,7 +561,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
for provider, label in self._iter_providers():
try:
result, error = await self._call_with_rate_limit(
@@ -574,7 +585,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
for provider, label in self._iter_providers():
try:
result = await self._call_with_rate_limit(
@@ -602,13 +613,14 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
for provider, label in self._iter_providers():
try:
result = await self._call_with_rate_limit(
label,
provider.get_user_models,
username,
cursor=cursor,
)
if result is not None:
return result
@@ -624,6 +636,19 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None
async def get_creator_model_count(self, username: str) -> Optional[int]:
for provider, label in self._iter_providers():
try:
result = await provider.get_creator_model_count(username)
if result is not None:
return result
except Exception as e:
logger.debug(
"Provider %s failed for get_creator_model_count: %s", label, e
)
continue
return None
def _iter_providers(self):
return zip(self.providers, self._provider_labels)
@@ -656,14 +681,14 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
def __getattr__(self, item):
return getattr(self._provider, item)
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_by_hash,
model_hash,
)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_versions,
@@ -673,7 +698,7 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
async def get_model_versions_bulk(
self,
model_ids: Sequence[int],
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_versions_bulk,
@@ -682,14 +707,14 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_versions_by_hashes,
hashes,
)
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_version,
@@ -697,20 +722,24 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
version_id,
)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_version_info,
version_id,
)
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_user_models,
username,
cursor=cursor,
)
async def get_creator_model_count(self, username: str) -> Optional[int]:
return await self._provider.get_creator_model_count(username)
class ModelMetadataProviderManager:
"""Manager for selecting and using model metadata providers"""
@@ -733,12 +762,12 @@ class ModelMetadataProviderManager:
if is_default or self.default_provider is None:
self.default_provider = name
async def get_model_by_hash(self, model_hash: str, provider_name: str = None) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str, provider_name: Optional[str] = None) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by hash using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_by_hash(model_hash)
async def get_model_versions(self, model_id: str, provider_name: str = None) -> Optional[Dict]:
async def get_model_versions(self, model_id: str, provider_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Get model versions using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_versions(model_id)
@@ -746,8 +775,8 @@ class ModelMetadataProviderManager:
async def get_model_versions_bulk(
self,
model_ids: Sequence[int],
provider_name: str = None,
) -> Optional[Dict[int, Dict]]:
provider_name: Optional[str] = None,
) -> Optional[Dict[int, Dict[str, Any]]]:
"""Fetch model versions for multiple model ids when supported by provider."""
provider = self._get_provider(provider_name)
try:
@@ -755,12 +784,12 @@ class ModelMetadataProviderManager:
except NotImplementedError:
return None
async def get_model_version(self, model_id: int = None, version_id: int = None, provider_name: str = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None, provider_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Get specific model version using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_version(model_id, version_id)
async def get_model_version_info(self, version_id: str, provider_name: str = None) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str, provider_name: Optional[str] = None) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version info using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_version_info(version_id)
@@ -768,20 +797,30 @@ class ModelMetadataProviderManager:
async def get_model_versions_by_hashes(
self,
hashes: List[str],
provider_name: str = None,
) -> Optional[List[Dict]]:
provider_name: Optional[str] = None,
) -> Optional[List[Dict[str, Any]]]:
provider = self._get_provider(provider_name)
try:
return await provider.get_model_versions_by_hashes(hashes)
except NotImplementedError:
return None
async def get_user_models(self, username: str, provider_name: str = None) -> Optional[List[Dict]]:
"""Fetch models owned by the specified user"""
async def get_user_models(
self,
username: str,
provider_name: Optional[str] = None,
cursor: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
"""Fetch one page of models owned by the specified user"""
provider = self._get_provider(provider_name)
return await provider.get_user_models(username)
return await provider.get_user_models(username, cursor)
async def get_creator_model_count(self, username: str, provider_name: Optional[str] = None) -> Optional[int]:
"""Best-effort published model count for the specified user"""
provider = self._get_provider(provider_name)
return await provider.get_creator_model_count(username)
def _get_provider(self, provider_name: str = None) -> ModelMetadataProvider:
def _get_provider(self, provider_name: Optional[str] = None) -> ModelMetadataProvider:
"""Get provider by name or default provider"""
if provider_name:
if provider_name not in self.providers:
+13 -4
View File
@@ -12,6 +12,7 @@ from typing import (
Tuple,
Protocol,
Callable,
cast,
)
from ..utils.constants import NSFW_LEVELS
@@ -85,6 +86,7 @@ class SortParams:
key: str
order: str
seed: Optional[str] = None
@dataclass(frozen=True)
@@ -116,7 +118,7 @@ class ModelCacheRepository:
async def fetch_sorted(self, params: SortParams) -> List[Dict[str, Any]]:
"""Fetch cached data pre-sorted according to ``params``."""
cache = await self.get_cache()
return await cache.get_sorted_data(params.key, params.order)
return await cache.get_sorted_data(params.key, params.order, params.seed)
@staticmethod
def parse_sort(sort_by: str) -> SortParams:
@@ -132,10 +134,17 @@ class ModelCacheRepository:
sort_key = sort_by.strip().lower() or "name"
order = "asc"
if order not in ("asc", "desc"):
seed = None
if sort_key == "random":
# Random sort: the portion after ':' is the shuffle seed.
# A stable seed keeps paginated requests consistent; order is
# meaningless for a random shuffle.
seed = order if order and order not in ("asc", "desc") else None
order = "asc"
elif order not in ("asc", "desc"):
order = "asc"
return SortParams(key=sort_key, order=order)
return SortParams(key=sort_key, order=order, seed=seed)
class ModelFilterSet:
@@ -301,7 +310,7 @@ class ModelFilterSet:
else:
include_tags.add(normalized)
else:
include_tags = {tag.strip().lower() for tag in tag_filters if tag}
include_tags = {tag.strip().lower() for tag in cast(Iterable[Any], tag_filters) if tag}
if include_tags:
tag_logic = criteria.tag_logic.lower() if criteria.tag_logic else "any"
+289 -44
View File
@@ -5,16 +5,16 @@ import asyncio
import time
import shutil
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Set, Type, Union
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Set, Type, Union, cast
from ..utils.models import BaseModelMetadata
from ..utils.models import BaseModelMetadata, autov3_from_civitai_files
from ..config import config
from ..utils.file_utils import find_preview_file, get_preview_extension, calculate_sha256
from ..utils.file_utils import find_preview_file, get_preview_extension, calculate_sha256, calculate_autov3
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
@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
class CacheBuildResult:
"""Represents the outcome of scanning model files for cache building."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
hash_index: ModelHashIndex
tags_count: Dict[str, int]
excluded_models: List[str]
@@ -59,7 +59,7 @@ class ModelScanner:
lock = cls._get_lock()
async with lock:
if cls not in cls._instances:
cls._instances[cls] = cls()
cls._instances[cls] = cls() # pyright: ignore[reportCallIssue]
return cls._instances[cls]
def __init__(self, model_type: str, model_class: Type[BaseModelMetadata], file_extensions: Set[str], hash_index: Optional[ModelHashIndex] = None):
@@ -78,7 +78,8 @@ class ModelScanner:
self.model_type = model_type
self.model_class = model_class
self.file_extensions = file_extensions
self._cache = None
self._cache: Any = None
self._cache_version: int = 0
self._hash_index = hash_index or ModelHashIndex()
self._tags_count = {} # Dictionary to store tag counts
self._is_initializing = False # Flag to track initialization state
@@ -86,6 +87,7 @@ class ModelScanner:
self._persistent_cache = get_persistent_cache()
self._name_display_mode = self._resolve_name_display_mode()
self._cancel_requested = False # Flag for cancellation
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
try:
loop = asyncio.get_running_loop()
except RuntimeError:
@@ -97,6 +99,25 @@ class ModelScanner:
# Register this service
asyncio.create_task(self._register_service())
@property
def cache_version(self) -> int:
"""Monotonic version counter for the in-memory cache.
Every write path that mutates scanner cache state calls
:meth:`bump_cache_version`, so consumers (e.g. RecipeScanner) can
detect when a cached derivation of the raw data is stale. Reads never
bump.
"""
return self._cache_version
def bump_cache_version(self) -> None:
"""Invalidate derived caches by incrementing the cache version.
Public because external services (model lifecycle, route handlers)
rewrite scanner raw_data directly and must be able to invalidate it.
"""
self._cache_version += 1
def on_library_changed(self) -> None:
"""Reset caches when the active library changes."""
self._persistent_cache = get_persistent_cache()
@@ -106,6 +127,7 @@ class ModelScanner:
self._excluded_models = []
self._is_initializing = False
self._name_display_mode = self._resolve_name_display_mode()
self.bump_cache_version()
try:
loop = asyncio.get_running_loop()
@@ -182,7 +204,7 @@ class ModelScanner:
is_mapping = isinstance(source, Mapping)
def get_value(key: str, default: Any = None) -> Any:
if is_mapping:
if isinstance(source, Mapping):
return source.get(key, default)
sentinel = object()
@@ -225,6 +247,19 @@ class ModelScanner:
if not isinstance(notes, str):
notes = str(notes)
# AutoV3 three-state contract: absent key / None = "not checked yet",
# "" = "checked but unavailable" (never re-read the header), else the
# 12-char lowercase hex value. A metadata object already follows the
# contract and is passed through unchanged; a payload dict only carries
# an explicit checked state when the key is present.
if is_mapping:
if 'autov3' in source:
entry_autov3 = source['autov3'] or ''
else:
entry_autov3 = None
else:
entry_autov3 = get_value('autov3', None)
entry: Dict[str, Any] = {
'file_path': normalized_path,
# file_name is always stored WITHOUT extension (e.g. "OWSMianne_ANIMA_V1",
@@ -238,6 +273,7 @@ class ModelScanner:
'size': int(get_value('size', 0) or 0),
'modified': float(get_value('modified', 0.0) or 0.0),
'sha256': (get_value('sha256', '') or '').lower(),
'autov3': entry_autov3,
'base_model': get_value('base_model', '') or '',
'preview_url': preview_url,
'preview_nsfw_level': int(get_value('preview_nsfw_level', 0) or 0),
@@ -473,6 +509,13 @@ class ModelScanner:
if sha_value and path:
hash_index.add_entry(sha_value.lower(), path)
# Rebuild the AutoV3 index from the persisted autov3_index rows. These
# cover every known autov3 -> path mapping regardless of whether a
# sha256 row also exists for the same file.
for autov3_value, path in persisted.autov3_hash_rows:
if autov3_value and path:
hash_index.add_autov3(autov3_value.lower(), path)
tags_count: Dict[str, int] = {}
adjusted_raw_data: List[Dict[str, Any]] = []
for item in persisted.raw_data:
@@ -541,8 +584,30 @@ class ModelScanner:
'scanner_type': self.model_type,
'pageType': page_type
})
# Schedule the one-time AutoV3 backfill task (at most once per process)
# so entries loaded from a persisted snapshot that predates autov3 get
# their checked state computed in the background. The task never blocks
# or crashes the load path.
if not self._autov3_backfill_scheduled:
self._autov3_backfill_scheduled = True
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None
if loop is not None:
loop.create_task(self._run_autov3_backfill())
return True
async def _run_autov3_backfill(self) -> None:
"""Backfill autov3 for entries loaded from the persisted cache that lack it."""
try:
from ..services.autov3_backfill_service import Autov3BackfillService # lazy import (module created by another unit)
await Autov3BackfillService.get_instance().backfill(self)
except Exception as exc:
logger.warning("AutoV3 backfill failed: %s", exc)
async def _save_persistent_cache(self, scan_result: CacheBuildResult) -> None:
if not scan_result or not getattr(self, '_persistent_cache', None):
return
@@ -555,6 +620,7 @@ class ModelScanner:
return
hash_snapshot = self._build_hash_index_snapshot(scan_result.hash_index)
autov3_snapshot = self._build_autov3_index_snapshot(scan_result.hash_index)
loop = asyncio.get_event_loop()
try:
await loop.run_in_executor(
@@ -563,7 +629,8 @@ class ModelScanner:
self.model_type,
list(scan_result.raw_data),
hash_snapshot,
list(scan_result.excluded_models)
list(scan_result.excluded_models),
autov3_snapshot,
)
except Exception as exc:
logger.warning("%s Scanner: Failed to persist cache: %s", self.model_type.capitalize(), exc)
@@ -589,6 +656,20 @@ class ModelScanner:
bucket.append(path)
return snapshot
def _build_autov3_index_snapshot(self, hash_index: Optional[ModelHashIndex]) -> Dict[str, List[str]]:
"""Build the autov3 -> [paths] snapshot for the persisted cache."""
snapshot: Dict[str, List[str]] = {}
if not hash_index:
return snapshot
for autov3_value, path in hash_index.get_all_autov3().items():
if not autov3_value or not path:
continue
bucket = snapshot.setdefault(autov3_value.lower(), [])
if path not in bucket:
bucket.append(path)
return snapshot
async def _persist_current_cache(self) -> None:
if self._cache is None or not getattr(self, '_persistent_cache', None):
return
@@ -712,7 +793,7 @@ class ModelScanner:
else:
await self._reconcile_cache()
return self._cache
return cast(ModelCache, self._cache)
async def _initialize_cache(self) -> None:
"""Initialize or refresh the cache"""
@@ -872,6 +953,8 @@ class ModelScanner:
)
continue
model_data = validation_result.entry
if model_data is None:
continue
self._ensure_license_flags(model_data)
# Add to cache
@@ -880,7 +963,11 @@ class ModelScanner:
# Update hash index if available
if 'sha256' in model_data and 'file_path' in model_data:
self._hash_index.add_entry(model_data['sha256'].lower(), model_data['file_path'])
self._hash_index.add_entry(
model_data['sha256'].lower(),
model_data['file_path'],
model_data.get('autov3') or None
)
# Update tags count
if 'tags' in model_data and model_data['tags']:
@@ -927,6 +1014,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[str] = set()
deduped: list[Dict[str, Any]] = []
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
@@ -945,6 +1051,7 @@ class ModelScanner:
logger.error(f"{self.model_type.capitalize()} Scanner: Error reconciling cache: {e}", exc_info=True)
finally:
self._is_initializing = False # Unset flag
self.bump_cache_version()
def is_initializing(self) -> bool:
"""Check if the scanner is currently initializing"""
@@ -1025,7 +1132,7 @@ class ModelScanner:
*,
hash_index: Optional[ModelHashIndex] = None,
excluded_models: Optional[List[str]] = None
) -> Dict:
) -> Optional[Dict[str, Any]]:
"""Process a single model file and return its metadata"""
hash_index = hash_index or self._hash_index
excluded_models = excluded_models if excluded_models is not None else self._excluded_models
@@ -1049,7 +1156,7 @@ class ModelScanner:
file_name = os.path.splitext(os.path.basename(file_path))[0]
file_info['name'] = file_name
metadata = self.model_class.from_civitai_info(version_info, file_info, file_path)
metadata = cast(Any, self.model_class).from_civitai_info(version_info, file_info, file_path)
metadata.preview_url = find_preview_file(file_name, os.path.dirname(file_path))
await MetadataManager.save_metadata(file_path, metadata)
logger.info(f"Created metadata from .civitai.info for {file_path} (Reason: .civitai.info was found but .metadata.json was missing)")
@@ -1086,6 +1193,8 @@ class ModelScanner:
if metadata is None:
metadata = await self._create_default_metadata(file_path)
assert metadata is not None
# Hook: allow subclasses to adjust metadata
metadata = self.adjust_metadata(metadata, file_path, root_path)
@@ -1111,6 +1220,36 @@ class ModelScanner:
except Exception as e:
logger.error(f"Failed to compute SHA256 for {file_path}: {e}")
# AutoV3 resolution: prefer the Civitai AutoV3 reported for the file
# whose SHA256 matches (authoritative for recipe matching), falling
# back to the embedded safetensors header hash only for models never
# checked before (autov3 is None). A checked-unavailable state ('')
# is only upgraded by Civitai data — the header is never re-read.
current_autov3 = model_data.get('autov3')
if current_autov3 in (None, ''):
try:
civitai_data = None
if isinstance(metadata, BaseModelMetadata):
civitai_data = metadata.civitai
elif isinstance(metadata, dict):
civitai_data = metadata.get("civitai")
autov3 = autov3_from_civitai_files(
civitai_data, model_data.get("sha256") or ""
) or ""
if not autov3 and current_autov3 is None:
autov3 = (calculate_autov3(os.path.realpath(file_path)) or '').lower()
if autov3 != current_autov3:
model_data['autov3'] = autov3
if isinstance(metadata, BaseModelMetadata):
metadata.autov3 = autov3
await MetadataManager.save_metadata(file_path, metadata)
elif isinstance(metadata, dict):
# Dict payload: JSON null encodes the checked-unavailable state.
metadata['autov3'] = autov3 or None
await MetadataManager.save_metadata(file_path, metadata)
except Exception as e:
logger.error(f"Failed to resolve AutoV3 for {file_path}: {e}")
# Skip excluded models
if model_data.get('exclude', False):
excluded_models.append(model_data['file_path'])
@@ -1150,6 +1289,8 @@ class ModelScanner:
self._log_duplicate_filename_summary()
self.bump_cache_version()
def _log_duplicate_filename_summary(self) -> None:
"""Log a batched summary of duplicate filename conflicts once per scan."""
# Duplicate filename detection is only relevant for LoRAs, which use
@@ -1183,7 +1324,7 @@ class ModelScanner:
async def _sync_download_history(
self,
raw_data: List[Mapping[str, Any]],
raw_data: Sequence[Mapping[str, Any]],
*,
source: str,
) -> None:
@@ -1232,7 +1373,7 @@ class ModelScanner:
) -> CacheBuildResult:
"""Collect metadata for all model files."""
raw_data: List[Dict] = []
raw_data: List[Dict[str, Any]] = []
hash_index = ModelHashIndex()
tags_count: Dict[str, int] = {}
excluded_models: List[str] = []
@@ -1296,6 +1437,8 @@ class ModelScanner:
)
continue
result = validation_result.entry
if result is None:
continue
self._ensure_license_flags(result)
raw_data.append(result)
@@ -1303,7 +1446,7 @@ class ModelScanner:
sha_value = result.get('sha256')
model_path = result.get('file_path')
if sha_value and model_path:
hash_index.add_entry(sha_value.lower(), model_path)
hash_index.add_entry(sha_value.lower(), model_path, result.get('autov3') or None)
for tag in result.get('tags') or []:
tags_count[tag] = tags_count.get(tag, 0) + 1
@@ -1335,7 +1478,7 @@ class ModelScanner:
excluded_models=excluded_models
)
async def add_model_to_cache(self, metadata_dict: Dict, folder: str = '') -> bool:
async def add_model_to_cache(self, metadata_dict: Dict[str, Any], folder: str = '') -> bool:
"""Add a model to the cache
Args:
@@ -1348,31 +1491,44 @@ class ModelScanner:
try:
if self._cache is None:
await self.get_cached_data()
assert self._cache is not None
# 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'])
self._hash_index.add_entry(
metadata_dict['sha256'],
metadata_dict['file_path'],
metadata_dict.get('autov3') or None,
)
await self._persist_current_cache()
self.bump_cache_version()
return True
except Exception as e:
logger.error(f"Error adding model to cache: {e}")
return False
async def move_model(self, source_path: str, target_path: str) -> Optional[str]:
async def move_model(self, source_path: str, target_path: str) -> Optional[Dict[str, Any]]:
"""Move a model and its associated files to a new location
Args:
@@ -1394,6 +1550,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)
@@ -1403,7 +1562,7 @@ class ModelScanner:
# Check for filename conflicts and auto-rename if necessary
from ..utils.models import BaseModelMetadata
final_filename = BaseModelMetadata.generate_unique_filename(
target_path, base_name, file_ext, get_source_hash
target_path, base_name, file_ext, lambda: get_source_hash() or ""
)
target_file = os.path.join(target_path, final_filename).replace(os.sep, '/')
@@ -1451,7 +1610,7 @@ class ModelScanner:
logger.error(f"Error moving associated file {source_file}: {e}")
# Handle metadata file specially to update paths
if source_metadata and os.path.exists(source_metadata):
if source_metadata and moved_metadata_path and os.path.exists(source_metadata):
try:
shutil.move(source_metadata, moved_metadata_path)
metadata = await self._update_metadata_paths(moved_metadata_path, target_file)
@@ -1469,7 +1628,7 @@ class ModelScanner:
logger.error(f"Error moving model: {e}", exc_info=True)
return None
async def _update_metadata_paths(self, metadata_path: str, model_path: str) -> Dict:
async def _update_metadata_paths(self, metadata_path: str, model_path: str) -> Optional[Dict[str, Any]]:
"""Update file paths in metadata file"""
try:
with open(metadata_path, 'r', encoding='utf-8') as f:
@@ -1495,7 +1654,7 @@ class ModelScanner:
logger.error(f"Error updating metadata paths: {e}", exc_info=True)
return None
async def update_single_model_cache(self, original_path: str, new_path: str, metadata: Dict, recalculate_type: bool = False) -> Union[bool, Dict]:
async def update_single_model_cache(self, original_path: str, new_path: str, metadata: Optional[Dict[str, Any]], recalculate_type: bool = False) -> Union[bool, Dict[str, Any]]:
"""Update cache after a model has been moved or modified"""
cache = await self.get_cached_data()
@@ -1518,6 +1677,7 @@ class ModelScanner:
]
cache_modified = bool(existing_item) or bool(metadata)
cache_entry: Optional[Dict[str, Any]] = None
if metadata:
normalized_new_path = new_path.replace(os.sep, '/')
@@ -1549,7 +1709,11 @@ class ModelScanner:
sha_value = cache_entry.get('sha256')
if sha_value:
self._hash_index.add_entry(sha_value.lower(), normalized_new_path)
self._hash_index.add_entry(
sha_value.lower(),
normalized_new_path,
cache_entry.get('autov3') or None,
)
all_folders = set(item['folder'] for item in cache.raw_data)
cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
@@ -1563,8 +1727,11 @@ class ModelScanner:
if cache_modified:
await self._persist_current_cache()
self.bump_cache_version()
return cache_entry if metadata else True
if metadata and cache_entry is not None:
return cache_entry
return True
async def sync_cache_from_metadata(
self, file_path: str, metadata_dict: Dict[str, Any]
@@ -1687,10 +1854,11 @@ class ModelScanner:
# ---- In-place update of the cache entry ----
existing_entry.clear()
existing_entry.update(desired_entry)
self.bump_cache_version()
# ---- Incremental tag count update ----
new_tags: set = set(desired_entry.get("tags") or [])
old_tag_set: set = set(old_tags)
new_tags: set[str] = set(desired_entry.get("tags") or [])
old_tag_set: set[str] = set(old_tags)
for tag in old_tag_set - new_tags:
current = self._tags_count.get(tag, 0)
if current <= 1:
@@ -1707,7 +1875,11 @@ class ModelScanner:
if old_sha:
self._hash_index.remove_by_path(file_path)
if new_sha:
self._hash_index.add_entry(new_sha, file_path)
self._hash_index.add_entry(
new_sha,
file_path,
desired_entry.get('autov3') or None,
)
# ---- Incremental version index update ----
new_civitai = desired_entry.get("civitai")
@@ -1723,7 +1895,7 @@ class ModelScanner:
# ---- 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
sort_key: Optional[str] = _last[0] if _last[0] is not None else None
if sort_key == "name":
if (
old_model_name != desired_entry.get("model_name", "")
@@ -1758,6 +1930,75 @@ class ModelScanner:
return True
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
"""Persist an AutoV3 hash for a single model (single write path used by the backfill service).
Locates the in-memory cache entry by ``file_path`` and updates only its
``autov3`` field: the in-memory hash index, the SQLite snapshot via
:meth:`PersistentModelCache.update_single_model`, and the
``.metadata.json`` sidecar. sha256, tags, and every other field are
left untouched, so the persistent delta only ever differs in autov3.
Returns:
``True`` when the entry was found and updated, ``False`` otherwise.
Never raises failures are logged and swallowed.
"""
try:
if self._cache is None:
return False
entry = next(
(item for item in self._cache.raw_data if item.get('file_path') == file_path),
None,
)
if entry is None:
return False
# Normalize once so the memory entry, sidecar, and SQLite row agree.
autov3 = (autov3 or "").lower()
# Capture the pre-mutation state so update_single_model only sees
# an autov3 delta between old and new.
old_item = dict(entry)
entry['autov3'] = autov3 or ''
# Prefer add_entry when a sha256 is known so the sha256 and autov3
# maps stay in sync; fall back to an autov3-only registration.
sha_value = entry.get('sha256')
checked_autov3 = entry.get('autov3') or None
if sha_value:
self._hash_index.add_entry(sha_value.lower(), file_path, checked_autov3)
elif checked_autov3:
self._hash_index.add_autov3(checked_autov3, file_path)
persistent = getattr(self, '_persistent_cache', None)
if persistent is not None:
await asyncio.get_event_loop().run_in_executor(
None,
persistent.update_single_model,
model_type,
entry,
old_item,
)
# Sidecar write-back: JSON null encodes the checked-unavailable
# state. Skip silently when the sidecar does not exist.
metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
if os.path.exists(metadata_path):
with open(metadata_path, 'r', encoding='utf-8') as handle:
payload = json.load(handle)
if not isinstance(payload, dict):
payload = {}
payload['autov3'] = entry['autov3'] or None
await MetadataManager.save_metadata(metadata_path, payload)
self.bump_cache_version()
return True
except Exception as exc:
logger.warning("Failed to update AutoV3 for %s: %s", file_path, exc)
return False
@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.
@@ -1817,7 +2058,7 @@ class ModelScanner:
return None
async def get_top_tags(self, limit: int = 20) -> List[Dict[str, any]]:
async def get_top_tags(self, limit: int = 20) -> List[Dict[str, Any]]:
"""Get top tags sorted by count. If limit is 0, return all tags."""
await self.get_cached_data()
@@ -1833,7 +2074,7 @@ class ModelScanner:
async def search_tags(
self, query: str, limit: int = 50
) -> List[Dict[str, any]]:
) -> 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``
@@ -1856,7 +2097,7 @@ class ModelScanner:
return matched
return matched[:limit]
async def get_base_models(self, limit: int = 20) -> List[Dict[str, any]]:
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()
@@ -1937,7 +2178,7 @@ class ModelScanner:
await self._persist_current_cache()
return updated
async def bulk_delete_models(self, file_paths: List[str]) -> Dict:
async def bulk_delete_models(self, file_paths: List[str]) -> Dict[str, Any]:
"""Delete multiple models and update cache in a batch operation
Args:
@@ -1971,6 +2212,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)
@@ -2083,6 +2326,8 @@ class ModelScanner:
await self._persist_current_cache()
self.bump_cache_version()
return True
except Exception as e:
@@ -2133,7 +2378,7 @@ class ModelScanner:
logger.error(f"Error checking model version existence: {e}")
return False
async def get_model_versions_by_id(self, model_id: int) -> List[Dict]:
async def get_model_versions_by_id(self, model_id: int) -> List[Dict[str, Any]]:
"""Get all versions of a model by its ID
Args:
+6 -6
View File
@@ -6,13 +6,13 @@ logger = logging.getLogger(__name__)
class ModelServiceFactory:
"""Factory for managing model services and routes"""
_services: Dict[str, Type] = {}
_routes: Dict[str, Type] = {}
_services: Dict[str, Type[Any]] = {}
_routes: Dict[str, Type[Any]] = {}
_initialized_services: Dict[str, Any] = {}
_initialized_routes: Dict[str, Any] = {}
@classmethod
def register_model_type(cls, model_type: str, service_class: Type, route_class: Type):
def register_model_type(cls, model_type: str, service_class: Type[Any], route_class: Type[Any]):
"""Register a new model type with its service and route classes
Args:
@@ -24,7 +24,7 @@ class ModelServiceFactory:
cls._routes[model_type] = route_class
@classmethod
def get_service_class(cls, model_type: str) -> Type:
def get_service_class(cls, model_type: str) -> Type[Any]:
"""Get service class for a model type
Args:
@@ -41,7 +41,7 @@ class ModelServiceFactory:
return cls._services[model_type]
@classmethod
def get_route_class(cls, model_type: str) -> Type:
def get_route_class(cls, model_type: str) -> Type[Any]:
"""Get route class for a model type
Args:
@@ -87,7 +87,7 @@ class ModelServiceFactory:
logger.error(f"Failed to setup routes for {model_type}: {e}", exc_info=True)
@classmethod
def get_registered_types(cls) -> list:
def get_registered_types(cls) -> list[str]:
"""Get list of all registered model types
Returns:
+29 -24
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""Service for tracking remote model version updates."""
from __future__ import annotations
@@ -336,9 +340,9 @@ class ModelUpdateService:
return
try:
from .persistent_model_cache import get_persistent_cache
from .persistent_model_cache import PersistentModelCache
legacy_path = get_persistent_cache(self._library_name).get_database_path()
legacy_path = PersistentModelCache.get_default(self._library_name).get_database_path()
except Exception:
return
@@ -735,7 +739,7 @@ class ModelUpdateService:
)
results: Dict[int, ModelUpdateRecord] = {}
prefetched: Dict[int, Mapping] = {}
prefetched: Dict[int, Mapping[Any, Any]] = {}
fetch_targets: List[int] = []
if metadata_provider and local_versions:
@@ -834,7 +838,7 @@ class ModelUpdateService:
model_id: int,
version_ids: Sequence[int],
*,
version_info: Optional[Mapping] = None,
version_info: Optional[Mapping[str, Any]] = None,
) -> ModelUpdateRecord:
"""Persist a new set of in-library version identifiers."""
@@ -954,7 +958,11 @@ class ModelUpdateService:
records = self._get_records_bulk(model_type, normalized_ids)
return {
model_id: records.get(model_id).has_update(hide_early_access=hide_early_access) if records.get(model_id) else False
model_id: (
records[model_id].has_update(hide_early_access=hide_early_access)
if model_id in records
else False
)
for model_id in normalized_ids
}
@@ -980,7 +988,7 @@ class ModelUpdateService:
metadata_provider,
*,
force_refresh: bool = False,
prefetched_response: Optional[Mapping] = None,
prefetched_response: Optional[Mapping[str, Any]] = None,
all_local_version_ids: Optional[Sequence[int]] = None,
) -> Optional[ModelUpdateRecord]:
normalized_local = self._normalize_sequence(local_versions)
@@ -1010,7 +1018,7 @@ class ModelUpdateService:
fallback_attempted = False
fallback_error_message: Optional[str] = None
mark_model_as_ignored = False
response: Optional[Mapping] = None
response: Optional[Mapping[str, Any]] = None
if metadata_provider and should_fetch:
response = prefetched_response
if response is None:
@@ -1122,7 +1130,7 @@ class ModelUpdateService:
async def _enrich_version_entries(
self,
metadata_provider,
responses_by_model_id: Dict[int, Mapping],
responses_by_model_id: Dict[int, Mapping[Any, Any]],
) -> None:
"""Enrich version entries with ``usageControl`` via batch hash endpoint.
@@ -1151,7 +1159,7 @@ class ModelUpdateService:
all_hashes = list(version_ids_by_hash.keys())
BATCH_SIZE = 100
enrichment: Dict[int, Dict] = {}
enrichment: Dict[int, Dict[str, Any]] = {}
try:
for start in range(0, len(all_hashes), BATCH_SIZE):
batch = all_hashes[start : start + BATCH_SIZE]
@@ -1208,7 +1216,7 @@ class ModelUpdateService:
version["earlyAccessEndsAt"] = extra["earlyAccessEndsAt"]
@staticmethod
def _collect_hashes_from_response(response: Mapping) -> Dict[int, str]:
def _collect_hashes_from_response(response: Mapping[str, Any]) -> Dict[int, str]:
"""Extract ``{version_id: sha256}`` from a model-level API response.
Returns an empty dict if the response structure is unexpected.
@@ -1229,7 +1237,7 @@ class ModelUpdateService:
return result
@staticmethod
def _extract_sha256_from_version_entry(entry: Mapping) -> Optional[str]:
def _extract_sha256_from_version_entry(entry: Mapping[str, Any]) -> Optional[str]:
"""Return the SHA256 hash from the primary model file of a version entry."""
files = entry.get("files")
if not isinstance(files, list):
@@ -1253,22 +1261,19 @@ class ModelUpdateService:
self,
metadata_provider,
model_ids: Sequence[int],
) -> Dict[int, Mapping]:
) -> Dict[int, Mapping[Any, Any]]:
"""Fetch model metadata in batches of up to 100 ids."""
BATCH_SIZE = 100
normalized = self._normalize_sequence(model_ids)
if not normalized:
provider = metadata_provider
if not normalized or provider is None:
return {}
aggregated: Dict[int, Mapping] = {}
aggregated: Dict[int, Mapping[Any, Any]] = {}
total_ids = len(normalized)
total_batches = (total_ids + BATCH_SIZE - 1) // BATCH_SIZE
provider_name = (
metadata_provider.__class__.__name__
if metadata_provider is not None
else "unknown"
)
provider_name = provider.__class__.__name__
for batch_index, start in enumerate(range(0, total_ids, BATCH_SIZE), start=1):
chunk = normalized[start : start + BATCH_SIZE]
logger.info(
@@ -1279,7 +1284,7 @@ class ModelUpdateService:
provider_name,
)
try:
response = await metadata_provider.get_model_versions_bulk(chunk)
response = await provider.get_model_versions_bulk(chunk)
except RateLimitError:
raise
if response is None:
@@ -1356,7 +1361,7 @@ class ModelUpdateService:
model_type: Optional[str] = None,
model_id: Optional[int] = None,
last_checked_at: Optional[float] = None,
version_info: Optional[Mapping] = None,
version_info: Optional[Mapping[str, Any]] = None,
) -> ModelUpdateRecord:
local_set = set(normalized_local)
# When folder-filtering, also consider versions in other folders
@@ -1578,7 +1583,7 @@ class ModelUpdateService:
if not isinstance(files, Iterable):
return None
def parse_size(entry: Mapping) -> Optional[int]:
def parse_size(entry: Mapping[str, Any]) -> Optional[int]:
size_kb = entry.get("sizeKB")
if size_kb is None:
return None
@@ -1664,8 +1669,8 @@ class ModelUpdateService:
return {}
ids = list(model_ids)
status_rows: list = []
version_rows: list = []
status_rows: list[sqlite3.Row] = []
version_rows: list[sqlite3.Row] = []
with self._connect() as conn:
for start in range(0, len(ids), self._SQLITE_MAX_VARIABLES):
+159 -21
View File
@@ -3,8 +3,8 @@ import logging
import os
import sqlite3
import threading
from dataclasses import dataclass
from typing import Dict, List, Mapping, Optional, Sequence, Tuple
from dataclasses import dataclass, field
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
@@ -15,9 +15,10 @@ logger = logging.getLogger(__name__)
class PersistedCacheData:
"""Lightweight structure returned by the persistent cache."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
hash_rows: List[Tuple[str, str]]
excluded_models: List[str]
autov3_hash_rows: List[Tuple[str, str]] = field(default_factory=list)
DEFAULT_LICENSE_FLAGS = 127 # 127 (0b1111111) encodes default CivitAI permissions with all commercial modes enabled.
@@ -36,6 +37,7 @@ class PersistentModelCache:
"size",
"modified",
"sha256",
"autov3",
"base_model",
"preview_url",
"preview_nsfw_level",
@@ -68,8 +70,8 @@ class PersistentModelCache:
self._db_path = db_path or self._resolve_default_path(self._library_name)
self._db_lock = threading.Lock()
self._schema_initialized = False
directory = os.path.dirname(self._db_path)
try:
directory = os.path.dirname(self._db_path)
if directory:
os.makedirs(directory, exist_ok=True)
except Exception as exc: # pragma: no cover - defensive guard
@@ -118,6 +120,10 @@ class PersistentModelCache:
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
(model_type,),
).fetchall()
autov3_rows = conn.execute(
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
(model_type,),
).fetchall()
excluded = conn.execute(
"SELECT file_path FROM excluded_models WHERE model_type = ?",
(model_type,),
@@ -128,7 +134,7 @@ class PersistentModelCache:
logger.warning("Failed to load persisted cache for %s: %s", model_type, exc)
return None
raw_data: List[Dict] = []
raw_data: List[Dict[str, Any]] = []
for row in rows:
file_path: str = row["file_path"]
trained_words = []
@@ -139,7 +145,7 @@ class PersistentModelCache:
trained_words = []
creator_username = row["civitai_creator_username"]
civitai: Optional[Dict] = None
civitai: Optional[Dict[str, Any]] = None
civitai_has_data = any(
row[col] is not None
for col in ("civitai_id", "civitai_model_id", "civitai_model_type", "civitai_name")
@@ -191,6 +197,8 @@ class PersistentModelCache:
"hash_status": row["hash_status"] or "completed",
"hf_url": row["hf_url"] or "",
}
if row["autov3"] is not None:
item["autov3"] = (row["autov3"] or "").lower()
raw_data.append(item)
hash_pairs = [(entry["sha256"].lower(), entry["file_path"]) for entry in hash_rows if entry["sha256"]]
@@ -201,10 +209,21 @@ class PersistentModelCache:
if sha_value:
hash_pairs.append((sha_value.lower(), item["file_path"]))
excluded_paths = [row["file_path"] for row in excluded]
return PersistedCacheData(raw_data=raw_data, hash_rows=hash_pairs, excluded_models=excluded_paths)
autov3_pairs = [
(entry["autov3"].lower(), entry["file_path"])
for entry in autov3_rows
if entry["autov3"]
]
def save_cache(self, model_type: str, raw_data: Sequence[Dict], hash_index: Dict[str, List[str]], excluded_models: Sequence[str]) -> None:
excluded_paths = [row["file_path"] for row in excluded]
return PersistedCacheData(
raw_data=raw_data,
hash_rows=hash_pairs,
excluded_models=excluded_paths,
autov3_hash_rows=autov3_pairs,
)
def save_cache(self, model_type: str, raw_data: Sequence[Dict[str, Any]], hash_index: Dict[str, List[str]], excluded_models: Sequence[str], autov3_hash_index: Optional[Dict[str, List[str]]] = None) -> None:
if not self.is_enabled():
return
if not self._schema_initialized:
@@ -219,7 +238,7 @@ class PersistentModelCache:
conn.execute("BEGIN")
model_rows = [self._prepare_model_row(model_type, item) for item in raw_data]
model_map: Dict[str, Tuple] = {
model_map: Dict[str, Tuple[Any, ...]] = {
row[1]: row for row in model_rows if row[1] # row[1] is file_path
}
@@ -251,13 +270,17 @@ class PersistentModelCache:
"DELETE FROM hash_index WHERE model_type = ? AND file_path = ?",
to_remove_models,
)
conn.executemany(
"DELETE FROM autov3_index WHERE model_type = ? AND file_path = ?",
to_remove_models,
)
conn.executemany(
"DELETE FROM excluded_models WHERE model_type = ? AND file_path = ?",
to_remove_models,
)
insert_rows: List[Tuple] = []
update_rows: List[Tuple] = []
insert_rows: List[Tuple[Any, ...]] = []
update_rows: List[Tuple[Any, ...]] = []
for file_path, row in model_map.items():
existing = existing_model_map.get(file_path)
@@ -289,11 +312,11 @@ class PersistentModelCache:
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
(model_type,),
).fetchall()
existing_tags: Dict[str, set] = {}
existing_tags: Dict[str, set[str]] = {}
for row in existing_tags_rows:
existing_tags.setdefault(row["file_path"], set()).add(row["tag"])
new_tags: Dict[str, set] = {}
new_tags: Dict[str, set[str]] = {}
for item in raw_data:
file_path = item.get("file_path")
if not file_path:
@@ -332,14 +355,14 @@ class PersistentModelCache:
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
(model_type,),
).fetchall()
existing_hash_map: Dict[str, set] = {}
existing_hash_map: Dict[str, set[str]] = {}
for row in existing_hash_rows:
sha_value = (row["sha256"] or "").lower()
if not sha_value:
continue
existing_hash_map.setdefault(sha_value, set()).add(row["file_path"])
new_hash_map: Dict[str, set] = {}
new_hash_map: Dict[str, set[str]] = {}
for sha_value, paths in hash_index.items():
normalized_sha = (sha_value or "").lower()
if not normalized_sha:
@@ -373,6 +396,52 @@ class PersistentModelCache:
hash_inserts,
)
if autov3_hash_index is not None:
existing_autov3_rows = conn.execute(
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
(model_type,),
).fetchall()
existing_autov3_map: Dict[str, set[str]] = {}
for row in existing_autov3_rows:
autov3_value = (row["autov3"] or "").lower()
if not autov3_value:
continue
existing_autov3_map.setdefault(autov3_value, set()).add(row["file_path"])
new_autov3_map: Dict[str, set[str]] = {}
for autov3_value, paths in autov3_hash_index.items():
normalized_autov3 = (autov3_value or "").lower()
if not normalized_autov3:
continue
bucket = new_autov3_map.setdefault(normalized_autov3, set())
for path in paths:
if path:
bucket.add(path)
autov3_inserts: List[Tuple[str, str, str]] = []
autov3_deletes: List[Tuple[str, str, str]] = []
all_autov3 = set(existing_autov3_map.keys()) | set(new_autov3_map.keys())
for autov3_value in all_autov3:
existing_paths = existing_autov3_map.get(autov3_value, set())
new_paths = new_autov3_map.get(autov3_value, set())
for path in existing_paths - new_paths:
autov3_deletes.append((model_type, autov3_value, path))
for path in new_paths - existing_paths:
autov3_inserts.append((model_type, autov3_value, path))
if autov3_deletes:
conn.executemany(
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
autov3_deletes,
)
if autov3_inserts:
conn.executemany(
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
autov3_inserts,
)
existing_excluded_rows = conn.execute(
"SELECT file_path FROM excluded_models WHERE model_type = ?",
(model_type,),
@@ -435,6 +504,7 @@ class PersistentModelCache:
size INTEGER,
modified REAL,
sha256 TEXT,
autov3 TEXT,
base_model TEXT,
preview_url TEXT,
preview_nsfw_level INTEGER,
@@ -472,6 +542,13 @@ class PersistentModelCache:
PRIMARY KEY (model_type, sha256, file_path)
);
CREATE TABLE IF NOT EXISTS autov3_index (
model_type TEXT NOT NULL,
autov3 TEXT NOT NULL,
file_path TEXT NOT NULL,
PRIMARY KEY (model_type, autov3, file_path)
);
CREATE TABLE IF NOT EXISTS excluded_models (
model_type TEXT NOT NULL,
file_path TEXT NOT NULL,
@@ -504,6 +581,7 @@ class PersistentModelCache:
"license_flags": f"INTEGER DEFAULT {DEFAULT_LICENSE_FLAGS}",
"hash_status": "TEXT DEFAULT 'completed'",
"hf_url": "TEXT DEFAULT ''",
"autov3": "TEXT",
}
for column, definition in required_columns.items():
@@ -522,7 +600,7 @@ class PersistentModelCache:
conn.row_factory = sqlite3.Row
return conn
def _prepare_model_row(self, model_type: str, item: Dict) -> Tuple:
def _prepare_model_row(self, model_type: str, item: Dict[str, Any]) -> Tuple[Any, ...]:
civitai = item.get("civitai") or {}
trained_words = civitai.get("trainedWords")
if isinstance(trained_words, str):
@@ -549,6 +627,12 @@ class PersistentModelCache:
if license_flags is None:
license_flags = DEFAULT_LICENSE_FLAGS
autov3_value = item.get("autov3")
if autov3_value is None:
autov3_column = None
else:
autov3_column = (autov3_value or "").lower()
return (
model_type,
item.get("file_path"),
@@ -558,6 +642,7 @@ class PersistentModelCache:
int(item.get("size") or 0),
float(item.get("modified") or 0.0),
(item.get("sha256") or "").lower() or None,
autov3_column,
item.get("base_model") or "",
item.get("preview_url") or "",
int(item.get("preview_nsfw_level") or 0),
@@ -590,8 +675,8 @@ class PersistentModelCache:
def update_single_model(
self,
model_type: str,
new_item: Dict,
old_item: Optional[Dict] = None,
new_item: Dict[str, Any],
old_item: Optional[Dict[str, Any]] = None,
) -> None:
"""Update a single model row in the persistent cache.
@@ -630,8 +715,8 @@ class PersistentModelCache:
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()
new_tags: set[str] = set(new_item.get("tags") or [])
old_tags: set[str] = 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
@@ -663,6 +748,25 @@ class PersistentModelCache:
(model_type, new_sha, file_path),
)
# --- autov3_index ---
new_autov3: Optional[str] = new_item.get("autov3")
if new_autov3 is not None:
new_autov3 = (new_autov3 or "").lower()
old_autov3: Optional[str] = (old_item.get("autov3") if old_item else None)
if old_autov3 is not None:
old_autov3 = (old_autov3 or "").lower()
if new_autov3 != old_autov3:
if old_autov3:
conn.execute(
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
(model_type, old_autov3, file_path),
)
if new_autov3:
conn.execute(
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
(model_type, new_autov3, file_path),
)
conn.execute("COMMIT")
except Exception:
conn.execute("ROLLBACK")
@@ -676,6 +780,40 @@ class PersistentModelCache:
exc,
)
def get_models_missing_autov3(self, model_type: str) -> List[str]:
"""Return file paths whose models lack an AutoV3 checked state.
Only rows with a completed sha256 and a NULL autov3 column qualify
rows with '' (checked-unavailable) or a value are never returned, so
the backfill query self-terminates.
"""
if not self.is_enabled():
return []
if not self._schema_initialized:
self._initialize_schema()
if not self._schema_initialized:
return []
try:
with self._db_lock:
conn = self._connect(readonly=True)
try:
rows = conn.execute(
"SELECT file_path FROM models "
"WHERE model_type = ? AND autov3 IS NULL "
"AND sha256 IS NOT NULL AND sha256 != ''",
(model_type,),
).fetchall()
finally:
conn.close()
return [row["file_path"] for row in rows]
except Exception as exc:
logger.warning(
"Failed to query models missing autov3 for %s: %s",
model_type,
exc,
)
return []
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 = ?",
+12 -8
View File
@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""SQLite-based persistent cache for recipe metadata.
This module provides fast recipe cache persistence using SQLite, enabling
@@ -13,7 +17,7 @@ import os
import sqlite3
import threading
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple
from typing import Any, Dict, List, Optional, Set, Tuple
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
@@ -24,7 +28,7 @@ logger = logging.getLogger(__name__)
class PersistedRecipeData:
"""Lightweight structure returned by the persistent recipe cache."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
file_stats: Dict[str, Tuple[float, int]] # json_path -> (mtime, size)
image_id_map: Dict[str, str] = field(default_factory=dict)
"""Precomputed mapping of civitai image_id → recipe_id."""
@@ -63,8 +67,8 @@ class PersistentRecipeCache:
self._db_path = db_path or self._resolve_default_path(self._library_name)
self._db_lock = threading.Lock()
self._schema_initialized = False
directory = os.path.dirname(self._db_path)
try:
directory = os.path.dirname(self._db_path)
if directory:
os.makedirs(directory, exist_ok=True)
except Exception as exc:
@@ -140,7 +144,7 @@ class PersistentRecipeCache:
logger.warning("Failed to load persisted recipe cache: %s", exc)
return None
raw_data: List[Dict] = []
raw_data: List[Dict[str, Any]] = []
file_stats: Dict[str, Tuple[float, int]] = {}
for row in rows:
@@ -162,7 +166,7 @@ class PersistentRecipeCache:
def save_cache(
self,
recipes: List[Dict],
recipes: List[Dict[str, Any]],
json_paths: Optional[Dict[str, str]] = None,
image_id_map: Optional[Dict[str, str]] = None,
) -> None:
@@ -251,7 +255,7 @@ class PersistentRecipeCache:
except Exception:
return {}
def update_recipe(self, recipe: Dict, json_path: Optional[str] = None) -> None:
def update_recipe(self, recipe: Dict[str, Any], json_path: Optional[str] = None) -> None:
"""Update or insert a single recipe in the cache.
Args:
@@ -439,7 +443,7 @@ class PersistentRecipeCache:
conn.row_factory = sqlite3.Row
return conn
def _prepare_recipe_row(self, recipe: Dict, json_path: str) -> Tuple:
def _prepare_recipe_row(self, recipe: Dict[str, Any], json_path: str) -> Tuple[Any, ...]:
"""Convert a recipe dict to a row tuple for SQLite insertion."""
loras = recipe.get("loras")
loras_json = json.dumps(loras) if loras else None
@@ -486,7 +490,7 @@ class PersistentRecipeCache:
tags_json,
)
def _row_to_recipe(self, row: sqlite3.Row) -> Dict:
def _row_to_recipe(self, row: sqlite3.Row) -> Dict[str, Any]:
"""Convert a SQLite row to a recipe dictionary."""
loras = []
if row["loras_json"]:
+3 -1
View File
@@ -22,7 +22,7 @@ class PreviewAssetService:
self,
*,
metadata_manager,
downloader_factory: Callable[[], Awaitable],
downloader_factory: Callable[[], Awaitable[Any]],
exif_utils,
) -> None:
self._metadata_manager = metadata_manager
@@ -69,6 +69,8 @@ class PreviewAssetService:
if not preview_url:
return
preview_url = str(preview_url)
def extension_from_url(url: str, fallback: str) -> str:
try:
parsed = urlparse(url)
+13 -12
View File
@@ -1,5 +1,5 @@
import asyncio
from typing import Iterable, List, Dict, Optional
from typing import Any, Iterable, List, Dict, Optional
from dataclasses import dataclass, field
from natsort import natsorted
@@ -8,12 +8,13 @@ from natsort import natsorted
class RecipeCache:
"""Cache structure for Recipe data"""
raw_data: List[Dict]
sorted_by_name: List[Dict]
sorted_by_date: List[Dict]
raw_data: List[Dict[str, Any]]
sorted_by_name: List[Dict[str, Any]]
sorted_by_date: List[Dict[str, Any]]
folders: List[str] | None = None
folder_tree: Dict | None = None
folder_tree: Dict[str, Any] | None = None
image_id_map: Dict[str, str] = field(default_factory=dict)
_lock: Any = field(init=False, repr=False, default=None)
"""Mapping of civitai image_id → recipe_id, precomputed at cache build time.
Built once during cache initialization (O(n)) so that
@@ -40,7 +41,7 @@ class RecipeCache:
)
async def update_recipe_metadata(
self, recipe_id: str, metadata: Dict, *, resort: bool = True
self, recipe_id: str, metadata: Dict[str, Any], *, resort: bool = True
) -> bool:
"""Update metadata for a specific recipe in all cached data
@@ -60,7 +61,7 @@ class RecipeCache:
return True
return False # Recipe not found
async def add_recipe(self, recipe_data: Dict, *, resort: bool = False) -> None:
async def add_recipe(self, recipe_data: Dict[str, Any], *, resort: bool = False) -> None:
"""Add a new recipe to the cache."""
async with self._lock:
@@ -70,7 +71,7 @@ class RecipeCache:
async def remove_recipe(
self, recipe_id: str, *, resort: bool = False
) -> Optional[Dict]:
) -> Optional[Dict[str, Any]]:
"""Remove a recipe from the cache by ID.
Args:
@@ -91,7 +92,7 @@ class RecipeCache:
async def bulk_remove(
self, recipe_ids: Iterable[str], *, resort: bool = False
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""Remove multiple recipes from the cache."""
id_set = {str(recipe_id) for recipe_id in recipe_ids}
@@ -111,7 +112,7 @@ class RecipeCache:
return removed
async def replace_recipe(
self, recipe_id: str, new_data: Dict, *, resort: bool = False
self, recipe_id: str, new_data: Dict[str, Any], *, resort: bool = False
) -> bool:
"""Replace cached data for a recipe."""
@@ -124,7 +125,7 @@ class RecipeCache:
return True
return False
async def get_recipe(self, recipe_id: str) -> Optional[Dict]:
async def get_recipe(self, recipe_id: str) -> Optional[Dict[str, Any]]:
"""Return a shallow copy of a cached recipe."""
async with self._lock:
@@ -133,7 +134,7 @@ class RecipeCache:
return dict(recipe)
return None
async def snapshot(self) -> List[Dict]:
async def snapshot(self) -> List[Dict[str, Any]]:
"""Return a copy of all cached recipes."""
async with self._lock:

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