diff --git a/CHANGELOG.md b/CHANGELOG.md index aa9265172..af5a8b9d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,7 @@ And check out new **history** tab in the right panel, it now shows visualization - configurable grid images size in *settings -> user interface -> grid image size* - gallery now includes reference model images + - reference models now include indicator if they are *ready* or *need download* - **Offloading** - **balanced** - enable offload during pre-forward by default diff --git a/modules/paths.py b/modules/paths.py index 85750a051..503f6eeb0 100644 --- a/modules/paths.py +++ b/modules/paths.py @@ -24,6 +24,7 @@ try: except Exception: config = {} +reference_path = os.path.join('models', 'Reference') modules_path = os.path.dirname(os.path.realpath(__file__)) script_path = os.path.dirname(modules_path) data_path = cli.data_dir diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 843c248cc..645dab049 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -300,8 +300,9 @@ class ExtraNetworksPage: self.html = '' self.create_items(tabname) versions = sorted({item.get("version", "") for item in self.items if item.get("version")}) - if 'ref' in versions: - versions.remove('ref') + for v in ['ref', 'reference', 'ready', 'download']: + if v in versions: + versions.remove(v) versions_html = '' for ver in versions: versions_html += f'
' diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py index a8c35c2c4..1d320f760 100644 --- a/modules/ui_extra_networks_checkpoints.py +++ b/modules/ui_extra_networks_checkpoints.py @@ -2,10 +2,9 @@ import os import html import json import concurrent -from modules import shared, ui_extra_networks, sd_models, modelstats +from modules import shared, ui_extra_networks, sd_models, modelstats, paths -reference_dir = os.path.join('models', 'Reference') version_map = { "QwenEdit": "Qwen", "Flux.1 D": "Flux", @@ -25,6 +24,12 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): shared.refresh_checkpoints() def list_reference(self): # pylint: disable=inconsistent-return-statements + existing = [model.filename if model.type == 'safetensors' else model.name for model in sd_models.checkpoints_list.values()] + + def reference_downloaded(url): + url = url.split('@')[0] if '@' in url else 'Diffusers/' + url + return any(model.endswith(url) for model in existing) + if not shared.opts.sd_checkpoint_autodownload or not shared.opts.extra_network_reference_enable: return [] for k, v in shared.reference_models.items(): @@ -36,15 +41,15 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): else: continue preview = v.get('preview', v['path']) - preview_file = self.find_preview_file(os.path.join(reference_dir, preview)) + preview_file = self.find_preview_file(os.path.join(paths.reference_path, preview)) _size, mtime = modelstats.stat(preview_file) - name = os.path.normpath(os.path.join(reference_dir, k)).replace('\\', '/') + name = os.path.normpath(os.path.join(paths.reference_path, k)).replace('\\', '/') yield { "type": 'Model', "name": name, "title": name, "filename": url, - "preview": self.find_preview(os.path.join(reference_dir, preview)), + "preview": self.find_preview(os.path.join(paths.reference_path, preview)), "local_preview": preview_file, "onclick": '"' + html.escape(f"selectReference({json.dumps(url)})") + '"', "hash": None, @@ -53,7 +58,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): "info": {}, "metadata": {}, "description": v.get('desc', ''), - "version": "ref", + "version": "ready" if reference_downloaded(url) else "download", } def create_item(self, name): @@ -101,4 +106,4 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): return items def allowed_directories_for_previews(self): - return [v for v in [shared.opts.ckpt_dir, reference_dir, sd_models.model_path] if v is not None] + return [v for v in [shared.opts.ckpt_dir, paths.reference_path, sd_models.model_path] if v is not None]