interrogate update

CHANGELOG.md

@@ -2,22 +2,26 @@
 
 ## Update for 2024-09-15
 
-- hf force logout/login on token change  
-- flux avoid unet load if unchanged  
-- flux mark specific unet as unavailable if load failed  
-- xyz grid full refactor  
-- xyz grid multi-mode: *selectable-script* and *alwayson-script*  
-- xyz grid allow usage combined with other scripts  
-- xyz grid allow **unet** selection  
-- xyz grid allow passing **model args** directly:  
-  allowed params will be checked against models call signature  
-  example: `width=768; height=512, width=512; height=768`  
-- xyz grid allow passing **processing args** directly:  
-  params are set directly on main processing object and can be known or new params  
-  example: `steps=10, steps=20; test=unknown`  
-- backend=original is now marked as in maintenance-only mode  
-- fix: minor ui optimizations  
-- fix: diffusers local model name parsing  
+- **flux**
+  - avoid unet load if unchanged  
+  - mark specific unet as unavailable if load failed  
+  - fix diffusers local model name parsing  
+- **xyz grid** full refactor  
+  - multi-mode: *selectable-script* and *alwayson-script*  
+  - allow usage combined with other scripts  
+  - allow **unet** selection  
+  - allow passing **model args** directly:  
+    allowed params will be checked against models call signature  
+    example: `width=768; height=512, width=512; height=768`  
+  - allow passing **processing args** directly:  
+    params are set directly on main processing object and can be known or new params  
+    example: `steps=10, steps=20; test=unknown`  
+- **interrogate**  
+  - add additional blip models: *blip-base, blip-large, blip-t5-xl, blip-t5-xxl, opt-2.7b, opt-6.7b*  
+  - add advanced params  
+- **hf** force logout/login on token change  
+- **backend=original** is now marked as in maintenance-only mode  
+- minor ui optimizations  
 
 ## Update for 2024-09-13
 

javascript/sdnext.css

@@ -36,7 +36,7 @@ input[type='color'] { width: 64px; height: 32px; }
 .gradio-button.secondary-down { background: var(--button-secondary-background-fill); color: var(--button-secondary-text-color); }
 .gradio-button.secondary-down, .gradio-button.secondary-down:hover { box-shadow: 1px 1px 1px rgba(0,0,0,0.25) inset, 0px 0px 3px rgba(0,0,0,0.15) inset; }
 .gradio-button.secondary-down:hover { background: var(--button-secondary-background-fill-hover); color: var(--button-secondary-text-color-hover); }
-.gradio-button.tool { max-width: min-content; min-width: min-content !important; align-self: end; font-size: 20px !important; color: var(--body-text-color) !important; align-self: end; }
+.gradio-button.tool { max-width: min-content; min-width: min-content !important; font-size: 20px !important; color: var(--body-text-color) !important; align-self: end; margin-bottom: 4px; }
 .gradio-checkbox { margin: 0.75em 1.5em 0 0; align-self: center; }
 .gradio-column { min-width: min(160px, 100%) !important; }
 .gradio-container { max-width: unset !important; padding: var(--block-label-padding) !important; }

modules/interrogate.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import time
 from collections import namedtuple
 from pathlib import Path
 import re
@@ -11,6 +12,24 @@ from torchvision.transforms.functional import InterpolationMode
 from modules import devices, paths, shared, lowvram, errors
 
 
+config = {
+    "caption_max_length": 64,
+    "chunk_size": 1024,
+    "flavor_intermediate_count": 1024,
+    "min_flavors": 2,
+    "max_flavors": 8,
+    "clip_offload": True,
+    "caption_offload": True,
+}
+caption_models = {
+    'blip-base': 'Salesforce/blip-image-captioning-base',
+    'blip-large': 'Salesforce/blip-image-captioning-large',
+    'blip2-opt-2.7b': 'Salesforce/blip2-opt-2.7b-coco',
+    'blip2-opt-6.7b': 'Salesforce/blip2-opt-6.7b',
+    'blip2-flip-t5-xl': 'Salesforce/blip2-flan-t5-xl',
+    'blip2-flip-t5-xxl': 'Salesforce/blip2-flan-t5-xxl',
+}
+ci = None
 blip_image_eval_size = 384
 clip_model_name = 'ViT-L/14'
 Category = namedtuple("Category", ["name", "topn", "items"])
@@ -200,10 +219,6 @@ class InterrogateModels:
 
 # --------- interrrogate ui
 
-ci = None
-low_vram = False
-
-
 class BatchWriter:
     def __init__(self, folder):
         self.folder = folder
@@ -219,24 +234,54 @@ class BatchWriter:
             self.file.close()
 
 
+def update_interrogate_params(caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count):
+    config["caption_max_length"] = int(caption_max_length)
+    config["chunk_size"] = int(chunk_size)
+    config["min_flavors"] = int(min_flavors)
+    config["max_flavors"] = int(max_flavors)
+    config["flavor_intermediate_count"] = int(flavor_intermediate_count)
+    if ci is not None:
+        ci.config.caption_max_length = config["caption_max_length"]
+        ci.config.chunk_size = config["chunk_size"]
+        ci.config.flavor_intermediate_count = config["flavor_intermediate_count"]
+    shared.log.debug(f'Interrogate params: {config}')
+
 def get_clip_models():
     import open_clip
     return ['/'.join(x) for x in open_clip.list_pretrained()]
 
 
-def load_interrogator(model):
-    from clip_interrogator import Config, Interrogator
+def load_interrogator(clip_model, blip_model):
+    import clip_interrogator
+    clip_interrogator.CAPTION_MODELS = caption_models
     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, 'Interrogator'), clip_model_name=model, quiet=True)
-        if low_vram:
-            config.apply_low_vram_defaults()
-        shared.log.info(f'Interrogate load: config={config}')
-        ci = Interrogator(config)
-    elif model != ci.config.clip_model_name:
-        ci.config.clip_model_name = model
-        shared.log.info(f'Interrogate load: config={ci.config}')
+        interrogator_config = clip_interrogator.Config(
+            device=devices.get_optimal_device(),
+            cache_path=os.path.join(paths.models_path, 'Interrogator'),
+            clip_model_name=clip_model,
+            caption_model_name=blip_model,
+            quiet=True,
+            caption_max_length=config['caption_max_length'],
+            chunk_size=config['chunk_size'],
+            flavor_intermediate_count=config['flavor_intermediate_count'],
+            clip_offload=config['clip_offload'],
+            caption_offload=config['caption_offload'],
+        )
+        t0 = time.time()
+        ci = clip_interrogator.Interrogator(interrogator_config)
+        t1 = time.time()
+        shared.log.info(f'Interrogate load: config={ci.config} min_flavors={config["min_flavors"]} max_flavors={config["max_flavors"]} time={t1-t0:.2f}')
+    elif clip_model != ci.config.clip_model_name or blip_model != ci.config.caption_model_name:
+        t0 = time.time()
+        ci.config.clip_model_name = clip_model
+        ci.config.clip_model = None
         ci.load_clip_model()
+        ci.config.caption_model_name = blip_model
+        ci.config.caption_model = None
+        ci.load_caption_model()
+        t1 = time.time()
+        shared.log.info(f'Interrogate reload: config={ci.config} min_flavors={config["min_flavors"]} max_flavors={config["max_flavors"]} time={t1-t0:.2f}')
 
 
 def unload_clip_model():
@@ -250,32 +295,35 @@ def unload_clip_model():
 
 
 def interrogate(image, mode, caption=None):
-    shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
+    shared.log.info(f'Interrogate: mode={mode} image={image}')
+    t0 = time.time()
     if mode == 'best':
-        prompt = ci.interrogate(image, caption=caption)
+        prompt = ci.interrogate(image, caption=caption, min_flavors=config["min_flavors"], max_flavors=config["max_flavors"])
     elif mode == 'caption':
         prompt = ci.generate_caption(image) if caption is None else caption
     elif mode == 'classic':
-        prompt = ci.interrogate_classic(image, caption=caption)
+        prompt = ci.interrogate_classic(image, caption=caption, max_flavors=config["max_flavors"])
     elif mode == 'fast':
-        prompt = ci.interrogate_fast(image, caption=caption)
+        prompt = ci.interrogate_fast(image, caption=caption, max_flavors=config["max_flavors"])
     elif mode == 'negative':
-        prompt = ci.interrogate_negative(image)
+        prompt = ci.interrogate_negative(image, max_flavors=config["max_flavors"])
     else:
         raise RuntimeError(f"Unknown mode {mode}")
+    t1 = time.time()
+    shared.log.debug(f'Interrogate: prompt="{prompt}" time={t1-t0:.2f}')
     return prompt
 
 
-def interrogate_image(image, model, mode):
+def interrogate_image(image, clip_model, blip_model, mode):
     shared.state.begin('Interrogate')
     try:
         if not shared.native and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
             lowvram.send_everything_to_cpu()
             devices.torch_gc()
-        load_interrogator(model)
+        load_interrogator(clip_model, blip_model)
         image = image.convert('RGB')
-        shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
         prompt = interrogate(image, mode)
+        devices.torch_gc()
     except Exception as e:
         prompt = f"Exception {type(e)}"
         shared.log.error(f'Interrogate: {e}')
@@ -283,7 +331,7 @@ def interrogate_image(image, model, mode):
     return prompt
 
 
-def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
+def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_model, mode, write):
     files = []
     if batch_files is not None:
         files += [f.name for f in batch_files]
@@ -300,7 +348,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
         if not shared.native and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
             lowvram.send_everything_to_cpu()
             devices.torch_gc()
-        load_interrogator(model)
+        load_interrogator(clip_model, blip_model)
         shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
         captions = []
         # first pass: generate captions
@@ -339,8 +387,8 @@ def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
     return '\n\n'.join(prompts)
 
 
-def analyze_image(image, model):
-    load_interrogator(model)
+def analyze_image(image, clip_model, blip_model):
+    load_interrogator(clip_model, blip_model)
     image = image.convert('RGB')
     image_features = ci.image_to_features(image)
     top_mediums = ci.mediums.rank(image_features, 5)

modules/sd_models.py

@@ -1247,7 +1247,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     if shared.opts.disable_accelerate:
                         from diffusers.utils import import_utils
                         import_utils._accelerate_available = False # pylint: disable=protected-access
-                    if shared.opts.diffusers_to_gpu:
+                    if shared.opts.diffusers_to_gpu and model_type.startswith('Stable Diffusion'):
+                        shared.log.debug(f'Diffusers accelerate: hijack={shared.opts.diffusers_to_gpu}')
                         sd_hijack_accelerate.hijack_accelerate()
                     else:
                         sd_hijack_accelerate.restore_accelerate()

modules/ui_postprocessing.py

@@ -45,17 +45,30 @@ def create_ui():
                         trending = gr.Label(elem_id="interrogate_label_trending", label="Trending", num_top_classes=5)
                         flavor = gr.Label(elem_id="interrogate_label_flavor", label="Flavor", num_top_classes=5)
                     with gr.Row():
-                        clip_model = gr.Dropdown([], value='ViT-L-14/openai', label='CLiP Model')
+                        clip_model = gr.Dropdown([], value='ViT-L-14/openai', label='CLiP model')
                         ui_common.create_refresh_button(clip_model, interrogate.get_clip_models, lambda: {"choices": interrogate.get_clip_models()}, 'refresh_interrogate_models')
-                        mode = gr.Radio(['best', 'fast', 'classic', 'caption', 'negative'], label='Mode', value='best')
+                        blip_model = gr.Dropdown(list(interrogate.caption_models), value='blip-base', label='Caption model')
+                        mode = gr.Dropdown(['best', 'fast', 'classic', 'caption', 'negative'], label='Mode', value='fast')
+                    with gr.Accordion(label='Advanced', open=False, visible=True):
+                        with gr.Row():
+                            caption_max_length = gr.Number(label='Max length', value=64, minimum=16, maximum=512, min_width=300)
+                            chunk_size = gr.Number(label='Chunk size', value=1024, minimum=256, maximum=4096, min_width=300)
+                            min_flavors = gr.Number(label='Min flavors', value=2, minimum=1, maximum=16, min_width=300)
+                            max_flavors = gr.Number(label='Max flavors', value=8, minimum=1, maximum=64, min_width=300)
+                            flavor_intermediate_count = gr.Number(label='Intermediates', value=1024, minimum=256, maximum=4096)
+                            caption_max_length.change(fn=interrogate.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
+                            chunk_size.change(fn=interrogate.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
+                            min_flavors.change(fn=interrogate.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
+                            max_flavors.change(fn=interrogate.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
+                            flavor_intermediate_count.change(fn=interrogate.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
                     with gr.Row(elem_id='interrogate_buttons_image'):
                         btn_interrogate_img = gr.Button("Interrogate", elem_id="interrogate_btn_interrogate", variant='primary')
                         btn_analyze_img = gr.Button("Analyze", elem_id="interrogate_btn_analyze", variant='primary')
                         btn_unload = gr.Button("Unload", elem_id="interrogate_btn_unload")
                     with gr.Row(elem_id='copy_buttons_interrogate'):
                         copy_interrogate_buttons = generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "extras", "control"])
-                    btn_interrogate_img.click(interrogate.interrogate_image, inputs=[image, clip_model, mode], outputs=prompt)
-                    btn_analyze_img.click(interrogate.analyze_image, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
+                    btn_interrogate_img.click(interrogate.interrogate_image, inputs=[image, clip_model, blip_model, mode], outputs=prompt)
+                    btn_analyze_img.click(interrogate.analyze_image, inputs=[image, clip_model, blip_model], outputs=[medium, artist, movement, trending, flavor])
                     btn_unload.click(interrogate.unload_clip_model)
                 with gr.Tab("Interrogate Batch"):
                     with gr.Row():
@@ -136,7 +149,7 @@ def create_ui():
     )
     btn_interrogate_batch.click(
         fn=interrogate.interrogate_batch,
-        inputs=[batch_files, batch_folder, batch_str, clip_model, mode, save_output],
+        inputs=[batch_files, batch_folder, batch_str, clip_model, blip_model, mode, save_output],
         outputs=[batch],
     )