Merge remote-tracking branch 'upstream/dev' into Extended-Merging

modules/images.py

@@ -413,7 +413,7 @@ class FilenameGenerator:
                 [part := part.replace(word, '_') for word in invalid_files] # pylint: disable=expression-not-assigned
             newparts.append(part)
         fn = Path(*newparts)
-        max_length = os.statvfs(__file__).f_namemax - 32 if hasattr(os, 'statvfs') else 230
+        max_length = max(230, os.statvfs(__file__).f_namemax - 32 if hasattr(os, 'statvfs') else 230)
         fn = str(fn)[:max_length-max(4, len(ext))].rstrip(invalid_suffix) + ext
         debug(f'Filename sanitize: input="{filename}" parts={parts} output="{fn}" ext={ext} max={max_length} len={len(fn)}')
         return fn

modules/loader.py

@@ -41,4 +41,4 @@ errors.install([gradio])
 
 import diffusers # pylint: disable=W0611,C0411
 timer.startup.record("diffusers")
-errors.log.debug(f'Load packages: torch={getattr(torch, "__long_version__", torch.__version__)} diffusers={diffusers.__version__} gradio={gradio.__version__}')
+errors.log.info(f'Load packages: torch={getattr(torch, "__long_version__", torch.__version__)} diffusers={diffusers.__version__} gradio={gradio.__version__}')

modules/modelloader.py

@@ -213,21 +213,24 @@ def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config
     pipeline_dir = None
 
     ok = True
+    err = None
     try:
         pipeline_dir = DiffusionPipeline.download(hub_id, **download_config)
     except Exception as e:
+        err = e
         ok = False
-        shared.log.warning(f"Diffusers download error: {hub_id} {e}")
-    if not ok:
+        # shared.log.warning(f"Diffusers download error: {hub_id} {e}")
+    if not ok and 'Repository Not Found' not in str(err):
         try:
             download_config.pop('load_connected_pipeline')
             download_config.pop('variant')
             pipeline_dir = hf.snapshot_download(hub_id, **download_config)
-        except Exception as e:
-            shared.log.warning(f"Diffusers hub download error: {hub_id} {e}")
+        except Exception:
+            # shared.log.warning(f"Diffusers download error: {hub_id} {e}")
+            pass
 
     if pipeline_dir is None:
-        shared.log.error(f"Diffusers no pipeline folder: {hub_id}")
+        shared.log.error(f"Diffusers download error: {hub_id} {err}")
         return None
     try:
         # TODO diffusers is this real error?
@@ -314,6 +317,23 @@ def find_diffuser(name: str):
     return None
 
 
+def load_reference(name: str):
+    found = [r for r in diffuser_repos if name == r['name'] or name == r['friendly'] or name == r['path']]
+    if len(found) > 0: # already downloaded
+        shared.log.debug(f'Reference model: {found[0]}')
+        return True
+    shared.log.debug(f'Reference download: {name}')
+    model_dir = download_diffusers_model(name, shared.opts.diffusers_dir)
+    if model_dir is None:
+        shared.log.debug(f'Reference download failed: {name}')
+        return False
+    else:
+        shared.log.debug(f'Reference download complete: {name}')
+        from modules import sd_models
+        sd_models.list_models()
+        return True
+
+
 modelloader_directories = {}
 cache_last = 0
 cache_time = 1

modules/processing.py

@@ -4,6 +4,7 @@ import math
 import time
 import hashlib
 import random
+import warnings
 from contextlib import nullcontext
 from typing import Any, Dict, List
 import torch
@@ -726,23 +727,16 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
 
 def validate_sample(sample):
-    ok = True
-    try:
-        sample = sample.astype(np.uint8)
-        return sample
-    except (Exception, Warning, RuntimeWarning) as e:
-        shared.log.error(f'Failed to validate sample values: {e}')
-        ok = False
-    if not ok:
-        try:
-            sample = np.nan_to_num(sample, nan=0, posinf=255, neginf=0)
-            sample = sample.astype(np.uint8)
-            shared.log.debug('Corrected sample values')
-        except (Exception, Warning, RuntimeWarning) as e:
-            shared.log.error(f'Failed to correct sample values: {e}')
-            sample = np.zeros_like(sample)
-            sample = sample.astype(np.uint8)
-    return sample
+    with warnings.catch_warnings(record=True) as w:
+        cast = sample.astype(np.uint8)
+    if len(w) > 0:
+        nans = np.isnan(sample).sum()
+        shared.log.error(f'Failed to validate samples: sample={sample.shape} invalid={nans}')
+        cast = np.nan_to_num(sample)
+        minimum, maximum, mean = np.min(cast), np.max(cast), np.mean(cast)
+        cast = cast.astype(np.uint8)
+        shared.log.warning(f'Attempted to correct samples: min={minimum:.2f} max={maximum:.2f} mean={mean:.2f}')
+    return cast
 
 
 def process_images_inner(p: StableDiffusionProcessing) -> Processed:

modules/processing_diffusers.py

@@ -143,6 +143,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro
             decoded = full_vae_decode(latents=latents, model=shared.sd_model)
         else:
             decoded = taesd_vae_decode(latents=latents)
+        # decoded = validate_sample(decoded) # TODO validate sample
         imgs = model.image_processor.postprocess(decoded, output_type=output_type)
         shared.state.job = prev_job
         return imgs

modules/sd_models.py

@@ -203,13 +203,13 @@ def list_models():
 def update_model_hashes():
     txt = []
     lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None]
-    shared.log.info(f'Models list: short hash missing for {len(lst)} out of {len(checkpoints_list)} models')
+    # shared.log.info(f'Models list: short hash missing for {len(lst)} out of {len(checkpoints_list)} models')
     for ckpt in lst:
         ckpt.hash = model_hash(ckpt.filename)
-        txt.append(f'Calculated short hash: <b>{ckpt.title}</b> {ckpt.hash}')
-    txt.append(f'Updated short hashes for <b>{len(lst)}</b> out of <b>{len(checkpoints_list)}</b> models')
+        # txt.append(f'Calculated short hash: <b>{ckpt.title}</b> {ckpt.hash}')
+    # txt.append(f'Updated short hashes for <b>{len(lst)}</b> out of <b>{len(checkpoints_list)}</b> models')
     lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.sha256 is None or ckpt.shorthash is None]
-    shared.log.info(f'Models list: full hash missing for {len(lst)} out of {len(checkpoints_list)} models')
+    shared.log.info(f'Models list: hash missing={len(lst)} total={len(checkpoints_list)}')
     for ckpt in lst:
         ckpt.sha256 = hashes.sha256(ckpt.filename, f"checkpoint/{ckpt.name}")
         ckpt.shorthash = ckpt.sha256[0:10] if ckpt.sha256 is not None else None

modules/ui.py

@@ -10,7 +10,7 @@ import numpy as np
 from PIL import Image
 from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call
 
-from modules import sd_hijack, sd_models, script_callbacks, ui_extensions, deepbooru, extra_networks, ui_common, ui_postprocessing, ui_loadsave, ui_train, ui_models
+from modules import sd_hijack, sd_models, script_callbacks, ui_extensions, deepbooru, extra_networks, ui_common, ui_postprocessing, ui_loadsave, ui_train, ui_models, ui_interrogate
 from modules.ui_components import FormRow, FormGroup, ToolButton, FormHTML
 from modules.paths import script_path, data_path
 from modules.shared import opts, cmd_opts
@@ -263,7 +263,7 @@ def create_toprow(is_img2img):
                 pause = gr.Button('Pause', elem_id=f"{id_part}_pause")
                 pause.click(fn=lambda: modules.shared.state.pause(), _js='checkPaused', inputs=[], outputs=[])
             with gr.Row(elem_id=f"{id_part}_tools"):
-                button_paste = gr.Button(value='Restore', variant='secondary', elem_id="paste") # symbols.paste
+                button_paste = gr.Button(value='Restore', variant='secondary', elem_id=f"{id_part}_paste") # symbols.paste
                 button_clear = gr.Button(value='Clear', variant='secondary', elem_id=f"{id_part}_clear_prompt_btn") # symbols.clear
                 button_extra = gr.Button(value='Networks', variant='secondary', elem_id=f"{id_part}_extra_networks_btn") # symbols.networks
                 button_clear.click(fn=lambda *x: ['', ''], inputs=[prompt, negative_prompt], outputs=[prompt, negative_prompt], show_progress=False)
@@ -273,8 +273,7 @@ def create_toprow(is_img2img):
                 negative_token_counter = gr.HTML(value="<span>0/75</span>", elem_id=f"{id_part}_negative_token_counter", elem_classes=["token-counter"])
                 negative_token_button = gr.Button(visible=False, elem_id=f"{id_part}_negative_token_button")
             with gr.Row(elem_id=f"{id_part}_styles_row"):
-                # prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[style.name for style in modules.shared.prompt_styles.styles.values()], value=[], multiselect=True)
-                prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=['aaa'], value=[], multiselect=True)
+                prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[style.name for style in modules.shared.prompt_styles.styles.values()], value=[], multiselect=True)
                 prompt_styles_btn_refresh = ToolButton(symbols.refresh, elem_id=f"{id_part}_styles_refresh", visible=True)
                 prompt_styles_btn_refresh.click(fn=lambda: gr.update(choices=[style.name for style in modules.shared.prompt_styles.styles.values()]), inputs=[], outputs=[prompt_styles])
                 prompt_styles_btn_select = gr.Button('Select', elem_id=f"{id_part}_styles_select", visible=False)
@@ -636,7 +635,6 @@ def create_ui(startup_timer = None):
                         img2img_batch_inpaint_mask_dir = gr.Textbox(label="Inpaint batch mask directory", **modules.shared.hide_dirs, elem_id="img2img_batch_inpaint_mask_dir")
 
                     img2img_tabs = [tab_img2img, tab_sketch, tab_inpaint, tab_inpaint_color, tab_inpaint_upload, tab_batch]
-
                     for i, tab in enumerate(img2img_tabs):
                         tab.select(fn=lambda tabnum=i: tabnum, inputs=[], outputs=[img2img_selected_tab])
 
@@ -904,6 +902,11 @@ def create_ui(startup_timer = None):
         ui_models.create_ui()
         timer.startup.record("ui-models")
 
+    with gr.Blocks(analytics_enabled=False) as interrogate_interface:
+        ui_interrogate.create_ui()
+        timer.startup.record("ui-interrogate")
+
+
     def create_setting_component(key, is_quicksettings=False):
         def fun():
             return opts.data[key] if key in opts.data else opts.data_labels[key].default
@@ -1103,11 +1106,12 @@ def create_ui(startup_timer = None):
     timer.startup.record("ui-settings")
 
     interfaces = [
-        (txt2img_interface, "From Text", "txt2img"),
-        (img2img_interface, "From Image", "img2img"),
-        (extras_interface, "Process Image", "process"),
+        (txt2img_interface, "Text", "txt2img"),
+        (img2img_interface, "Image", "img2img"),
+        (extras_interface, "Process", "process"),
         (train_interface, "Train", "train"),
         (models_interface, "Models", "models"),
+        (interrogate_interface, "Interrogate", "interrogate"),
     ]
     interfaces += script_callbacks.ui_tabs_callback()
     interfaces += [(settings_interface, "System", "system")]
@@ -1153,9 +1157,9 @@ def create_ui(startup_timer = None):
             inputs=components,
             outputs=[text_settings, result],
         )
-        defaults_submit.click(fn=lambda: modules.shared.restore_defaults(restart=True), _js="restart_reload")
-        restart_submit.click(fn=lambda: modules.shared.restart_server(restart=True), _js="restart_reload")
-        shutdown_submit.click(fn=lambda: modules.shared.restart_server(restart=False), _js="restart_reload")
+        defaults_submit.click(fn=lambda: modules.shared.restore_defaults(restart=True), _js="restartReload")
+        restart_submit.click(fn=lambda: modules.shared.restart_server(restart=True), _js="restartReload")
+        shutdown_submit.click(fn=lambda: modules.shared.restart_server(restart=False), _js="restartReload")
 
         for _i, k, _item in quicksettings_list:
             component = component_dict[k]
@@ -1190,6 +1194,21 @@ def create_ui(startup_timer = None):
             outputs=[component_dict['sd_vae'], text_settings],
         )
 
+        def reference_submit(model):
+            from modules import modelloader
+            loaded = modelloader.load_reference(model)
+            if loaded:
+                return model if loaded else opts.sd_model_checkpoint
+            print('HERE', model, loaded)
+            return loaded
+
+        button_set_reference = gr.Button('Change reference', elem_id='change_reference', visible=False)
+        button_set_reference.click(
+            fn=reference_submit,
+            _js="function(v){ return desiredCheckpointName; }",
+            inputs=[component_dict['sd_model_checkpoint']],
+            outputs=[component_dict['sd_model_checkpoint']],
+        )
         component_keys = [k for k in opts.data_labels.keys() if k in component_dict]
 
         def get_settings_values():

modules/ui_extra_networks.py

@@ -15,19 +15,19 @@ from collections import OrderedDict
 import gradio as gr
 from PIL import Image
 from starlette.responses import FileResponse, JSONResponse
-from modules import shared, scripts, modelloader
+from modules import paths, shared, scripts, modelloader
 from modules.ui_components import ToolButton
 import modules.ui_symbols as symbols
 
+
 allowed_dirs = []
 dir_cache = {} # key=path, value=(mtime, listdir(path))
 refresh_time = 0
 extra_pages = shared.extra_networks
 debug = shared.log.info if os.environ.get('SD_EN_DEBUG', None) is not None else lambda *args, **kwargs: None
 card_full = '''
-    <div class='card' onclick={card_click} title='{name}' data-tab='{tabname}' data-page='{page}' data-name='{name}' data-filename='{filename}' data-tags='{tags}' data-mtime='{mtime}' data-size='{size}'>
+    <div class='card' onclick={card_click} title='{name}' data-tab='{tabname}' data-page='{page}' data-name='{name}' data-filename='{filename}' data-tags='{tags}' data-mtime='{mtime}' data-size='{size}' data-search='{search}'>
         <div class='overlay'>
-            <span style="display:none" class='search_term'>{search_term}</span>
             <div class='tags'></div>
             <div class='name'>{title}</div>
         </div>
@@ -39,22 +39,21 @@ card_full = '''
     </div>
 '''
 card_list = '''
-    <div class='card card-list' onclick={card_click} title='{name}' data-tab='{tabname}' data-page='{page}' data-name='{name}' data-filename='{filename}' data-tags='{tags}' data-mtime='{mtime}' data-size='{size}'>
+    <div class='card card-list' onclick={card_click} title='{name}' data-tab='{tabname}' data-page='{page}' data-name='{name}' data-filename='{filename}' data-tags='{tags}' data-mtime='{mtime}' data-size='{size}' data-search='{search}'>
         <span class='details' title="Get details" onclick="showCardDetails(event)">&#x1f6c8;</span>&nbsp;
         <div class='name'>{title}</div>&nbsp;
         <div class='tags tags-list'></div>
-        <span style="display:none" class='search_term'>{search_term}</span>
     </div>
 '''
 
 
 def listdir(path):
-    debug(f'EN list-dir: {path}')
     if not os.path.exists(path):
         return []
     if path in dir_cache and os.path.getmtime(path) == dir_cache[path][0]:
         return dir_cache[path][1]
     else:
+        # debug(f'EN list-dir list: {path}')
         dir_cache[path] = (os.path.getmtime(path), [os.path.join(path, f) for f in os.listdir(path)])
         return dir_cache[path][1]
 
@@ -138,6 +137,9 @@ class ExtraNetworksPage:
         self.refresh_time = 0
         self.page_time = 0
         self.list_time = 0
+        self.info_time = 0
+        self.desc_time = 0
+        self.dirs = {}
         self.view = shared.opts.extra_networks_view
         self.card = card_full if shared.opts.extra_networks_view == 'gallery' else card_list
 
@@ -210,7 +212,6 @@ class ExtraNetworksPage:
             self.missing_thumbs.clear()
 
     def create_items(self, tabname):
-        debug(f'EN create-items: {self.name}')
         if self.refresh_time is not None and self.refresh_time > refresh_time: # cached results
             return
         t0 = time.time()
@@ -223,7 +224,8 @@ class ExtraNetworksPage:
         for item in self.items:
             self.metadata[item["name"]] = item.get("metadata", {})
         t1 = time.time()
-        self.list_time = round(t1-t0, 2)
+        debug(f'EN create-items: page={self.name} items={len(self.items)} time={t1-t0:.2f}')
+        self.list_time += t1-t0
 
 
     def create_page(self, tabname, skip = False):
@@ -237,8 +239,11 @@ class ExtraNetworksPage:
         allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews()]
         for parentdir, dirs in {d: modelloader.directory_directories(d) for d in allowed_folders}.items():
             for tgt in dirs.keys():
-                if shared.opts.diffusers_dir in tgt:
-                    subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
+                if shared.backend == shared.Backend.DIFFUSERS:
+                    if os.path.join(paths.models_path, 'Reference') in tgt:
+                        subdirs['Reference'] = 1
+                    if shared.opts.diffusers_dir in tgt:
+                        subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
                 if 'models--' in tgt:
                     continue
                 subdir = tgt[len(parentdir):].replace("\\", "/")
@@ -255,6 +260,7 @@ class ExtraNetworksPage:
         self.create_items(tabname)
         self.create_xyz_grid()
         htmls = []
+        self.items.sort(key=lambda x: x["mtime"], reverse=True)
         for item in self.items:
             htmls.append(self.create_html(item, tabname))
         self.html += ''.join(htmls)
@@ -263,7 +269,7 @@ class ExtraNetworksPage:
             self.html = f"<div id='{tabname}_{self_name_id}_subdirs' class='extra-network-subdirs'>{subdirs_html}</div><div id='{tabname}_{self_name_id}_cards' class='extra-network-cards'>{self.html}</div>"
         else:
             return ''
-        shared.log.debug(f"Extra networks: page='{self.name}' items={len(self.items)} subdirs={len(subdirs)} tab={tabname} dirs={self.allowed_directories_for_previews()} time={self.list_time}")
+        shared.log.debug(f"Extra networks: page='{self.name}' items={len(self.items)} subdirs={len(subdirs)} tab={tabname} dirs={self.allowed_directories_for_previews()} list={self.list_time:.2f} desc={self.desc_time:.2f} info={self.info_time:.2f}")
         if len(self.missing_thumbs) > 0:
             threading.Thread(target=self.create_thumb).start()
         return self.html
@@ -280,7 +286,7 @@ class ExtraNetworksPage:
                 "tabname": tabname,
                 "page": self.name,
                 "name": item["name"],
-                "title": item["name"].replace('_', ' '),
+                "title": os.path.basename(item["name"].replace('_', ' ')),
                 "filename": item["filename"],
                 "tags": '|'.join([item.get("tags")] if isinstance(item.get("tags", {}), str) else list(item.get("tags", {}).keys())),
                 "preview": html.escape(item.get("preview", self.link_preview('html/card-no-preview.png'))),
@@ -288,7 +294,7 @@ class ExtraNetworksPage:
                 "height": shared.opts.extra_networks_card_size if shared.opts.extra_networks_card_square else 'auto',
                 "fit": shared.opts.extra_networks_card_fit,
                 "prompt": item.get("prompt", None),
-                "search_term": item.get("search_term", ""),
+                "search": item.get("search_term", ""),
                 "description": item.get("description") or "",
                 "card_click": item.get("onclick", '"' + html.escape(f'return cardClicked({item.get("prompt", None)}, {"true" if self.allow_negative_prompt else "false"})') + '"'),
                 "mtime": item.get("mtime", 0),
@@ -305,8 +311,9 @@ class ExtraNetworksPage:
     def find_preview_file(self, path):
         fn = os.path.splitext(path)[0]
         preview_extensions = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"]
+        files = listdir(os.path.dirname(path))
         for file in [f'{fn}{mid}{ext}' for ext in preview_extensions for mid in ['.thumb.', '.preview.', '.']]:
-            if os.path.exists(file):
+            if file in files:
                 return file
         return 'html/card-no-preview.png'
 
@@ -315,14 +322,16 @@ class ExtraNetworksPage:
             return self.link_preview('html/card-no-preview.png')
         fn = os.path.splitext(path)[0]
         preview_extensions = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"]
+        files = listdir(os.path.dirname(path))
         for file in [f'{fn}{mid}{ext}' for ext in preview_extensions for mid in ['.thumb.', '.', '.preview.']]:
-            if os.path.exists(file):
+            if file in files:
                 if '.thumb.' not in file:
                     self.missing_thumbs.append(file)
                 return self.link_preview(file)
         return self.link_preview('html/card-no-preview.png')
 
-    def find_description(self, path):
+    def find_description(self, path, info=None):
+        t0 = time.time()
         class HTMLFilter(HTMLParser):
             text = ""
             def handle_data(self, data):
@@ -332,7 +341,8 @@ class ExtraNetworksPage:
                     self.text += '\n'
 
         fn = os.path.splitext(path)[0] + '.txt'
-        if os.path.exists(fn):
+        # if os.path.exists(fn):
+        if fn in listdir(os.path.dirname(path)):
             try:
                 with open(fn, "r", encoding="utf-8", errors="replace") as f:
                     txt = f.read()
@@ -340,20 +350,27 @@ class ExtraNetworksPage:
                     return txt
             except OSError:
                 pass
-        info = self.find_info(path)
+        if info is None:
+            info = self.find_info(path)
         desc = info.get('description', '') or ''
         f = HTMLFilter()
         f.feed(desc)
+        t1 = time.time()
+        self.desc_time += t1-t0
         return f.text
 
     def find_info(self, path):
+        t0 = time.time()
         fn = os.path.splitext(path)[0] + '.json'
-        if os.path.exists(fn):
+        # if os.path.exists(fn):
+        data = {}
+        if fn in listdir(os.path.dirname(path)):
             data = shared.readfile(fn, silent=True)
             if type(data) is list:
                 data = data[0]
-            return data
-        return {}
+        t1 = time.time()
+        self.info_time += t1-t0
+        return data
 
 
 def initialize():
@@ -524,8 +541,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         for page in get_pages():
             page.create_page(ui.tabname, skip_indexing)
             with gr.Tab(page.title, id=page.title.lower().replace(" ", "_"), elem_classes="extra-networks-tab") as tab:
-                hmtl = gr.HTML(page.html, elem_id=f'{tabname}{page.name}_extra_page', elem_classes="extra-networks-page")
-                ui.pages.append(hmtl)
+                page_html = gr.HTML(page.html, elem_id=f'{tabname}{page.name}_extra_page', elem_classes="extra-networks-page")
+                ui.pages.append(page_html)
                 tab.select(ui_tab_change, _js="getENActivePage", inputs=[ui.button_details], outputs=[ui.button_scan, ui.button_save, ui.button_model])
         # ui.tabs.change(fn=ui_tab_change, inputs=[], outputs=[ui.button_scan, ui.button_save])
 
@@ -724,7 +741,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         return ui_refresh_click(title)
 
     def ui_save_click():
-        from modules import paths, generation_parameters_copypaste
+        from modules import generation_parameters_copypaste
         filename = os.path.join(paths.data_path, "params.txt")
         if os.path.exists(filename):
             with open(filename, "r", encoding="utf8") as file:
@@ -736,7 +753,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         return res
 
     def ui_quicksave_click(name):
-        from modules import paths, generation_parameters_copypaste
+        from modules import generation_parameters_copypaste
         fn = os.path.join(paths.data_path, "params.txt")
         if os.path.exists(fn):
             with open(fn, "r", encoding="utf8") as file:

modules/ui_extra_networks_checkpoints.py

@@ -1,8 +1,9 @@
 import html
 import json
 import os
-from modules import shared, ui_extra_networks, sd_models
+from modules import shared, ui_extra_networks, sd_models, paths
 
+reference_dir = os.path.join(paths.models_path, 'Reference')
 
 class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
     def __init__(self):
@@ -11,12 +12,35 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
     def refresh(self):
         shared.refresh_checkpoints()
 
+    def list_reference(self):
+        if shared.backend != shared.Backend.DIFFUSERS:
+            return []
+        reference_models = shared.readfile(os.path.join('html', 'reference.json'))
+        for k, v in reference_models.items():
+            name = os.path.join(reference_dir, k)
+            yield {
+                "type": 'Model',
+                "name": name,
+                "title": name,
+                "filename": v['path'],
+                "search_term": self.search_terms_from_path(name),
+                "preview": self.find_preview(os.path.join(reference_dir, os.path.basename(v['path']))),
+                "local_preview": f"{os.path.splitext(name)[0]}.{shared.opts.samples_format}",
+                "onclick": '"' + html.escape(f"""return selectReference({json.dumps(v['path'])})""") + '"',
+                "hash": None,
+                "mtime": 0,
+                "size": 0,
+                "info": {},
+                "metadata": {},
+                "description": v.get('desc', ''),
+            }
+
     def list_items(self):
         checkpoint: sd_models.CheckpointInfo
         checkpoints = sd_models.checkpoints_list.copy()
         for name, checkpoint in checkpoints.items():
             try:
-                fn = os.path.splitext(checkpoint.filename)[0]
+                exists = os.path.exists(checkpoint.filename)
                 record = {
                     "type": 'Model',
                     "name": checkpoint.name,
@@ -24,18 +48,20 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
                     "filename": checkpoint.filename,
                     "hash": checkpoint.shorthash,
                     "search_term": self.search_terms_from_path(checkpoint.title),
-                    "preview": self.find_preview(fn),
-                    "local_preview": f"{fn}.{shared.opts.samples_format}",
-                    "description": self.find_description(fn),
-                    "info": self.find_info(fn),
+                    "preview": self.find_preview(checkpoint.filename),
+                    "local_preview": f"{os.path.splitext(checkpoint.filename)[0]}.{shared.opts.samples_format}",
                     "metadata": checkpoint.metadata,
                     "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
-                    "mtime": os.path.getmtime(checkpoint.filename),
-                    "size": os.path.getsize(checkpoint.filename),
+                    "mtime": os.path.getmtime(checkpoint.filename) if exists else 0,
+                    "size": os.path.getsize(checkpoint.filename) if exists else 0,
                 }
+                record["info"] = self.find_info(checkpoint.filename)
+                record["description"] = self.find_description(checkpoint.filename, record["info"])
                 yield record
             except Exception as e:
                 shared.log.debug(f"Extra networks error: type=model file={name} {e}")
+        for record in self.list_reference():
+            yield record
 
     def allowed_directories_for_previews(self):
-        return [v for v in [shared.opts.ckpt_dir, shared.opts.diffusers_dir, sd_models.model_path] if v is not None]
+        return [v for v in [shared.opts.ckpt_dir, shared.opts.diffusers_dir, reference_dir, sd_models.model_path] if v is not None]

modules/ui_extra_networks_hypernets.py

@@ -14,17 +14,17 @@ class ExtraNetworksPageHypernetworks(ui_extra_networks.ExtraNetworksPage):
         for name, path in shared.hypernetworks.items():
             try:
                 fn = os.path.splitext(path)[0]
-                name = os.path.relpath(fn, shared.opts.hypernetwork_dir)
+                name = os.path.relpath(os.path.splitext(path)[0], shared.opts.hypernetwork_dir)
                 yield {
                     "type": 'Hypernetwork',
-                    "name": os.path.relpath(fn, shared.opts.hypernetwork_dir),
+                    "name": name,
                     "filename": path,
-                    "preview": self.find_preview(fn),
-                    "description": self.find_description(fn),
-                    "info": self.find_info(fn),
+                    "preview": self.find_preview(path),
+                    "description": self.find_description(path),
+                    "info": self.find_info(path),
                     "search_term": self.search_terms_from_path(name),
-                    "prompt": json.dumps(f"<hypernet:{name}:{shared.opts.extra_networks_default_multiplier}>"),
-                    "local_preview": f"{fn}.{shared.opts.samples_format}",
+                    "prompt": json.dumps(f"<hypernet:{os.path.basename(name)}:{shared.opts.extra_networks_default_multiplier}>"),
+                    "local_preview": f"{os.path.splitext(path)[0]}.{shared.opts.samples_format}",
                     "mtime": os.path.getmtime(path),
                     "size": os.path.getsize(path),
                 }

modules/ui_extra_networks_textual_inversion.py

@@ -45,20 +45,21 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
                 if embedding.tag is not None:
                     tags[embedding.tag]=1
                 name = os.path.splitext(embedding.basename)[0]
-                yield {
+                record = {
                     "type": 'Embedding',
                     "name": name,
                     "filename": embedding.filename,
-                    "preview": self.find_preview(path),
-                    "description": self.find_description(path),
-                    "info": self.find_info(path),
+                    "preview": self.find_preview(embedding.filename),
                     "search_term": self.search_terms_from_path(name),
-                    "prompt": json.dumps(os.path.splitext(embedding.name)[0]),
+                    "prompt": json.dumps(f" {os.path.splitext(embedding.name)[0]}"),
                     "local_preview": f"{path}.{shared.opts.samples_format}",
                     "tags": tags,
                     "mtime": os.path.getmtime(embedding.filename),
                     "size": os.path.getsize(embedding.filename),
                 }
+                record["info"] = self.find_info(embedding.filename)
+                record["description"] = self.find_description(embedding.filename, record["info"])
+                yield record
             except Exception as e:
                 shared.log.debug(f"Extra networks error: type=embedding file={embedding.filename} {e}")
 

modules/ui_extra_networks_vae.py

@@ -14,23 +14,22 @@ class ExtraNetworksPageVAEs(ui_extra_networks.ExtraNetworksPage):
     def list_items(self):
         for name, filename in sd_vae.vae_dict.items():
             try:
-                fn = os.path.splitext(filename)[0]
                 record = {
                     "type": 'VAE',
                     "name": name,
                     "title": name,
-                    "filename": fn,
-                    "hash": hashes.sha256_from_cache(filename, f"vae/{fn}"),
-                    "search_term": self.search_terms_from_path(fn),
-                    "preview": self.find_preview(fn),
-                    "local_preview": f"{fn}.{shared.opts.samples_format}",
-                    "description": self.find_description(fn),
-                    "info": self.find_info(fn),
+                    "filename": filename,
+                    "hash": hashes.sha256_from_cache(filename, f"vae/{filename}"),
+                    "search_term": self.search_terms_from_path(filename),
+                    "preview": self.find_preview(filename),
+                    "local_preview": f"{os.path.splitext(filename)[0]}.{shared.opts.samples_format}",
                     "metadata": {},
                     "onclick": '"' + html.escape(f"""return selectVAE({json.dumps(name)})""") + '"',
                     "mtime": os.path.getmtime(filename),
                     "size": os.path.getsize(filename),
                 }
+                record["info"] = self.find_info(filename)
+                record["description"] = self.find_description(filename, record["info"])
                 yield record
             except Exception as e:
                 shared.log.debug(f"Extra networks error: type=vae file={filename} {e}")

modules/ui_interrogate.py

@@ -0,0 +1,267 @@
+import os
+import base64
+from io import BytesIO
+import gradio as gr
+import open_clip
+import torch
+from PIL import Image
+from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
+from fastapi import FastAPI
+from fastapi.exceptions import HTTPException
+from clip_interrogator import Config, Interrogator
+import modules.generation_parameters_copypaste as parameters_copypaste
+from modules import devices, lowvram, shared, paths
+
+
+ci = None
+low_vram = False
+
+
+class BatchWriter:
+    def __init__(self, folder):
+        self.folder = folder
+        self.csv, self.file = None, None
+
+    def add(self, file, prompt):
+        txt_file = os.path.splitext(file)[0] + ".txt"
+        with open(os.path.join(self.folder, txt_file), 'w', encoding='utf-8') as f:
+            f.write(prompt)
+
+    def close(self):
+        if self.file is not None:
+            self.file.close()
+
+
+def load(clip_model_name):
+    global ci # pylint: disable=global-statement
+    if ci is None:
+        config = Config(device=devices.get_optimal_device(), cache_path=os.path.join(paths.models_path, 'clip-interrogator'), clip_model_name=clip_model_name, quiet=True)
+        if low_vram:
+            config.apply_low_vram_defaults()
+        shared.log.info(f'Interrogate load: config={config}')
+        ci = Interrogator(config)
+    elif clip_model_name != ci.config.clip_model_name:
+        ci.config.clip_model_name = clip_model_name
+        shared.log.info(f'Interrogate load: config={ci.config}')
+        ci.load_clip_model()
+
+
+def unload():
+    if ci is not None:
+        shared.log.debug('Interrogate offload')
+        ci.caption_model = ci.caption_model.to(devices.cpu)
+        ci.clip_model = ci.clip_model.to(devices.cpu)
+        ci.caption_offloaded = True
+        ci.clip_offloaded = True
+        devices.torch_gc()
+
+
+def image_analysis(image, clip_model_name):
+    load(clip_model_name)
+    image = image.convert('RGB')
+    image_features = ci.image_to_features(image)
+    top_mediums = ci.mediums.rank(image_features, 5)
+    top_artists = ci.artists.rank(image_features, 5)
+    top_movements = ci.movements.rank(image_features, 5)
+    top_trendings = ci.trendings.rank(image_features, 5)
+    top_flavors = ci.flavors.rank(image_features, 5)
+    medium_ranks = dict(zip(top_mediums, ci.similarities(image_features, top_mediums)))
+    artist_ranks = dict(zip(top_artists, ci.similarities(image_features, top_artists)))
+    movement_ranks = dict(zip(top_movements, ci.similarities(image_features, top_movements)))
+    trending_ranks = dict(zip(top_trendings, ci.similarities(image_features, top_trendings)))
+    flavor_ranks = dict(zip(top_flavors, ci.similarities(image_features, top_flavors)))
+    return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks
+
+
+def interrogate(image, mode, caption=None):
+    shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
+    if mode == 'best':
+        prompt = ci.interrogate(image, caption=caption)
+    elif mode == 'caption':
+        prompt = ci.generate_caption(image) if caption is None else caption
+    elif mode == 'classic':
+        prompt = ci.interrogate_classic(image, caption=caption)
+    elif mode == 'fast':
+        prompt = ci.interrogate_fast(image, caption=caption)
+    elif mode == 'negative':
+        prompt = ci.interrogate_negative(image)
+    else:
+        raise RuntimeError(f"Unknown mode {mode}")
+    return prompt
+
+
+def image_to_prompt(image, mode, clip_model_name):
+    shared.state.begin()
+    shared.state.job = 'interrogate'
+    try:
+        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+            lowvram.send_everything_to_cpu()
+            devices.torch_gc()
+        load(clip_model_name)
+        image = image.convert('RGB')
+        shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
+        prompt = interrogate(image, mode)
+    except Exception as e:
+        prompt = f"Exception {type(e)}"
+        shared.log.error(f'Interrogate: {e}')
+    shared.state.end()
+    return prompt
+
+
+def get_models():
+    return ['/'.join(x) for x in open_clip.list_pretrained()]
+
+
+def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write):
+    files = []
+    if batch_files is not None:
+        files += [f.name for f in batch_files]
+    if batch_folder is not None:
+        files += [f.name for f in batch_folder]
+    if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
+        files += [os.path.join(batch_str, f) for f in os.listdir(batch_str) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+    if len(files) == 0:
+        shared.log.error('Interrogate batch no images')
+        return ''
+    shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
+    shared.state.begin()
+    shared.state.job = 'batch interrogate'
+    prompts = []
+    try:
+        if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
+            lowvram.send_everything_to_cpu()
+            devices.torch_gc()
+        load(clip_model)
+        captions = []
+        # first pass: generate captions
+        for file in files:
+            caption = ""
+            try:
+                if shared.state.interrupted:
+                    break
+                image = Image.open(file).convert('RGB')
+                caption = ci.generate_caption(image)
+            except Exception as e:
+                shared.log.error(f'Interrogate caption: {e}')
+            finally:
+                captions.append(caption)
+        # second pass: interrogate
+        if write:
+            writer = BatchWriter(os.path.dirname(files[0]))
+        for idx, file in enumerate(files):
+            try:
+                if shared.state.interrupted:
+                    break
+                image = Image.open(file).convert('RGB')
+                prompt = interrogate(image, mode, caption=captions[idx])
+                prompts.append(prompt)
+                if write:
+                    writer.add(file, prompt)
+            except OSError as e:
+                shared.log.error(f'Interrogate batch: {e}')
+        if write:
+            writer.close()
+        ci.config.quiet = False
+        unload()
+    except Exception as e:
+        shared.log.error(f'Interrogate batch: {e}')
+    shared.state.end()
+    return '\n\n'.join(prompts)
+
+
+def create_ui():
+    global low_vram # pylint: disable=global-statement
+    low_vram = shared.cmd_opts.lowvram or shared.cmd_opts.medvram
+    if not low_vram and torch.cuda.is_available():
+        device = devices.get_optimal_device()
+        vram_total = torch.cuda.get_device_properties(device).total_memory
+        if vram_total <= 12*1024*1024*1024:
+            low_vram = True
+    with gr.Row(elem_id="interrogate_tab"):
+        with gr.Column():
+            with gr.Tab("Image"):
+                with gr.Row():
+                    image = gr.Image(type='pil', label="Image")
+                with gr.Row():
+                    prompt = gr.Textbox(label="Prompt", lines=3)
+                with gr.Row():
+                    medium = gr.Label(label="Medium", num_top_classes=5)
+                    artist = gr.Label(label="Artist", num_top_classes=5)
+                    movement = gr.Label(label="Movement", num_top_classes=5)
+                    trending = gr.Label(label="Trending", num_top_classes=5)
+                    flavor = gr.Label(label="Flavor", num_top_classes=5)
+                with gr.Row():
+                    interrogate_btn = gr.Button("Interrogate", variant='primary')
+                    analyze_btn = gr.Button("Analyze", variant='primary')
+                    unload_btn = gr.Button("Unload")
+                with gr.Row():
+                    buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "extras"])
+                for tabname, button in buttons.items():
+                    parameters_copypaste.register_paste_params_button(parameters_copypaste.ParamBinding(paste_button=button, tabname=tabname, source_text_component=prompt, source_image_component=image,))
+            with gr.Tab("Batch"):
+                with gr.Row():
+                    batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], type='file', interactive=True, height=100)
+                with gr.Row():
+                    batch_folder = gr.File(label="Folder", show_label=True, file_count='directory', file_types=['image'], type='file', interactive=True, height=100)
+                with gr.Row():
+                    batch_str = gr.Text(label="Folder", value="", interactive=True)
+                with gr.Row():
+                    batch = gr.Text(label="Prompts", lines=10)
+                with gr.Row():
+                    write = gr.Checkbox(label='Write prompts to files', value=False)
+                with gr.Row():
+                    batch_btn = gr.Button("Interrogate", variant='primary')
+        with gr.Column():
+            with gr.Row():
+                clip_model = gr.Dropdown(get_models(), value='ViT-L-14/openai', label='CLIP Model')
+            with gr.Row():
+                mode = gr.Radio(['best', 'fast', 'classic', 'caption', 'negative'], label='Mode', value='best')
+        interrogate_btn.click(image_to_prompt, inputs=[image, mode, clip_model], outputs=prompt)
+        analyze_btn.click(image_analysis, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
+        unload_btn.click(unload)
+        batch_btn.click(batch_process, inputs=[batch_files, batch_folder, batch_str, mode, clip_model, write], outputs=[batch])
+
+
+def decode_base64_to_image(encoding):
+    if encoding.startswith("data:image/"):
+        encoding = encoding.split(";")[1].split(",")[1]
+    try:
+        image = Image.open(BytesIO(base64.b64decode(encoding)))
+        return image
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="Invalid encoded image") from e
+
+
+def mount_interrogator_api(_: gr.Blocks, app: FastAPI): # TODO redesign interrogator api
+
+    class InterrogatorAnalyzeRequest(BaseModel):
+        image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
+        clip_model_name: str = Field(default="ViT-L-14/openai", title="Model", description="The interrogate model used. See the models endpoint for a list of available models.")
+
+    class InterrogatorPromptRequest(InterrogatorAnalyzeRequest):
+        mode: str = Field(default="fast", title="Mode", description="The mode used to generate the prompt. Can be one of: best, fast, classic, negative.")
+
+    @app.get("/interrogator/models")
+    async def api_get_models():
+        return ["/".join(x) for x in open_clip.list_pretrained()]
+
+    @app.post("/interrogator/prompt")
+    async def api_get_prompt(analyzereq: InterrogatorPromptRequest):
+        image_b64 = analyzereq.image
+        if image_b64 is None:
+            raise HTTPException(status_code=404, detail="Image not found")
+        img = decode_base64_to_image(image_b64)
+        prompt = image_to_prompt(img, analyzereq.mode, analyzereq.clip_model_name)
+        return {"prompt": prompt}
+
+    @app.post("/interrogator/analyze")
+    async def api_analyze(analyzereq: InterrogatorAnalyzeRequest):
+        image_b64 = analyzereq.image
+        if image_b64 is None:
+            raise HTTPException(status_code=404, detail="Image not found")
+
+        img = decode_base64_to_image(image_b64)
+        (medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks) = image_analysis(img, analyzereq.clip_model_name)
+        return {"medium": medium_ranks, "artist": artist_ranks, "movement": movement_ranks, "trending": trending_ranks, "flavor": flavor_ranks}
+
+# script_callbacks.on_app_started(mount_interrogator_api)

modules/ui_postprocessing.py

@@ -25,7 +25,7 @@ def create_ui():
                 with gr.TabItem('Single Image', id="single_image", elem_id="extras_single_tab") as tab_single:
                     extras_image = gr.Image(label="Source", source="upload", interactive=True, type="pil", elem_id="extras_image")
                 with gr.TabItem('Process Batch', id="batch_process", elem_id="extras_batch_process_tab") as tab_batch:
-                    image_batch = gr.Files(label="Batch Process", interactive=True, elem_id="extras_image_batch")
+                    image_batch = gr.Files(label="Batch process", interactive=True, elem_id="extras_image_batch")
                 with gr.TabItem('Process Folder', id="batch_from_directory", elem_id="extras_batch_directory_tab") as tab_batch_dir:
                     extras_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, placeholder="A directory on the same machine where the server is running.", elem_id="extras_batch_input_dir")
                     extras_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, placeholder="Leave blank to save images to the default path.", elem_id="extras_batch_output_dir")

modules/ui_tempdir.py

@@ -3,7 +3,7 @@ import tempfile
 from collections import namedtuple
 from pathlib import Path
 import gradio as gr
-from PIL import PngImagePlugin
+from PIL import Image, PngImagePlugin
 from modules import shared, errors
 
 
@@ -36,7 +36,7 @@ def check_tmp_file(gradio, filename):
     return ok
 
 
-def pil_to_temp_file(self, img, dir: str, format="png") -> str: # pylint: disable=redefined-builtin,unused-argument
+def pil_to_temp_file(self, img: Image, dir: str, format="png") -> str: # pylint: disable=redefined-builtin,unused-argument
     """
     # original gradio implementation
     bytes_data = gr.processing_utils.encode_pil_to_bytes(img, format)
@@ -62,9 +62,12 @@ def pil_to_temp_file(self, img, dir: str, format="png") -> str: # pylint: disabl
         if isinstance(key, str) and isinstance(value, str):
             metadata.add_text(key, value)
             use_metadata = True
+    if not os.path.exists(dir):
+        os.makedirs(dir, exist_ok=True)
+        shared.log.debug(f'Created temp folder: path="{dir}"')
     with tempfile.NamedTemporaryFile(delete=False, suffix=".png", dir=dir) as tmp:
-        img.save(tmp, pnginfo=(metadata if use_metadata else None))
         name = tmp.name
+        img.save(name, pnginfo=(metadata if use_metadata else None))
         shared.log.debug(f'Saving temp: image="{name}"')
     return name