From 5bbc8e4b03ba1c74d8b37e108f29f4c6b5411c06 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 15 Sep 2024 10:50:52 -0400 Subject: [PATCH] interrogate update --- CHANGELOG.md | 36 +++++++------ javascript/sdnext.css | 2 +- modules/interrogate.py | 100 ++++++++++++++++++++++++++--------- modules/sd_models.py | 3 +- modules/ui_postprocessing.py | 23 ++++++-- 5 files changed, 115 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b0eb61303..818fae2a1 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/javascript/sdnext.css b/javascript/sdnext.css index dcabfaf8a..cdf300a26 100644 --- a/javascript/sdnext.css +++ b/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; } diff --git a/modules/interrogate.py b/modules/interrogate.py index eed43c773..90334c64a 100644 --- a/modules/interrogate.py +++ b/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) diff --git a/modules/sd_models.py b/modules/sd_models.py index 6e061dbed..eb5c78706 100644 --- a/modules/sd_models.py +++ b/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() diff --git a/modules/ui_postprocessing.py b/modules/ui_postprocessing.py index aa5140da2..e9ac2d72a 100644 --- a/modules/ui_postprocessing.py +++ b/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], )