Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor interrogate/analyze/vqa code

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-02-01 11:47:20 -05:00
parent ceaf0238fa
commit 654f44f66f
14 changed files with 64 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+1
View File
@@ -30,6 +30,7 @@
> python modules/api/nvml.py
- **Refactor**:
- unified trace handler with configurable tracebacks
- refactor interrogate/analyze/vqa code
- **Fixes**:
- photomaker with offloading
- photomaker with refine
modules/api/endpoints.py
+6 -5
View File
@@ -74,7 +74,7 @@ def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, f
return res
def get_interrogate():
from modules.interrogate import get_clip_models
from modules.interrogate.legacy import get_clip_models
return ['clip', 'deepdanbooru'] + get_clip_models()
def post_interrogate(req: models.ReqInterrogate):
@@ -84,16 +84,17 @@ def post_interrogate(req: models.ReqInterrogate):
image = image.convert('RGB')
if req.model == "clip":
try:
caption = shared.interrogator.interrogate(image)
from modules.interrogate import legacy
caption = legacy.interrogator.interrogate(image)
except Exception as e:
caption = str(e)
return models.ResInterrogate(caption=caption)
elif req.model == "deepdanbooru" or req.model == 'deepbooru':
from modules import deepbooru
from modules.interrogate import deepbooru
caption = deepbooru.model.tag(image)
return models.ResInterrogate(caption=caption)
else:
from modules.interrogate import interrogate_image, analyze_image, get_clip_models
from modules.interrogate.legacy import interrogate_image, analyze_image, get_clip_models
if req.model not in get_clip_models():
raise HTTPException(status_code=404, detail="Model not found")
try:
@@ -111,7 +112,7 @@ def post_vqa(req: models.ReqVQA):
raise HTTPException(status_code=404, detail="Image not found")
image = helpers.decode_base64_to_image(req.image)
image = image.convert('RGB')
from modules import vqa
from modules.interrogate import vqa
answer = vqa.interrogate(req.question, image, req.model)
return models.ResVQA(answer=answer)
modules/deepbooru.py → modules/interrogate/deepbooru.py
+3 -2
View File
@@ -4,7 +4,7 @@ import threading
import torch
import numpy as np
from PIL import Image
from modules import modelloader, paths, deepbooru_model, devices, images, shared
from modules import modelloader, paths, devices, images, shared
re_special = re.compile(r'([\\()])')
load_lock = threading.Lock()
@@ -27,7 +27,8 @@ class DeepDanbooru:
download_name='model-resnet_custom_v3.pt',
)
self.model = deepbooru_model.DeepDanbooruModel()
from modules.interrogate.deepbooru_model import DeepDanbooruModel
self.model = DeepDanbooruModel()
self.model.load_state_dict(torch.load(files[0], map_location="cpu"))
self.model.eval()
modules/deepbooru_model.py → modules/interrogate/deepbooru_model.py
View File
modules/interrogate/interrogate.py
+4
View File
@@ -0,0 +1,4 @@
def interrogate(image):
from modules.interrogate import legacy
prompt = legacy.interrogator.interrogate(image)
return prompt
modules/interrogate.py → modules/interrogate/legacy.py
+9 -4
View File
@@ -10,7 +10,7 @@ import torch.hub # pylint: disable=ungrouped-imports
from PIL import Image
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
from modules import devices, paths, shared, lowvram, errors
from modules import devices, paths, shared, lowvram, errors, sd_models
config = {
@@ -39,7 +39,7 @@ load_lock = threading.Lock()
def category_types():
return [f.stem for f in Path(shared.interrogator.content_dir).glob('*.txt')]
return [f.stem for f in Path(interrogator.content_dir).glob('*.txt')]
def download_default_clip_interrogate_categories(content_dir):
@@ -65,10 +65,10 @@ class InterrogateModels:
dtype = None
running_on_cpu = None
def __init__(self, content_dir):
def __init__(self, content_dir: str = None):
self.loaded_categories = None
self.skip_categories = []
self.content_dir = content_dir
self.content_dir = content_dir or os.path.join(paths.models_path, "interrogate")
self.running_on_cpu = False
def categories(self):
@@ -327,6 +327,8 @@ def interrogate_image(image, clip_model, blip_model, mode):
if not shared.native and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
lowvram.send_everything_to_cpu()
devices.torch_gc()
if shared.native:
sd_models.apply_balanced_offload(shared.sd_model)
load_interrogator(clip_model, blip_model)
image = image.convert('RGB')
prompt = interrogate(image, mode)
@@ -410,3 +412,6 @@ def analyze_image(image, clip_model, blip_model):
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
interrogator = InterrogateModels()
modules/vqa.py → modules/interrogate/vqa.py
View File
modules/sd_offload.py
+6
View File
@@ -179,6 +179,12 @@ def apply_balanced_offload(sd_model, exclude=[]):
global offload_hook_instance # pylint: disable=global-statement
if shared.opts.diffusers_offload_mode != "balanced":
return sd_model
if sd_model is None:
if not shared.sd_loaded:
return sd_model
sd_model = shared.sd_model
if sd_model is None:
return sd_model
t0 = time.time()
excluded = ['OmniGenPipeline']
if sd_model.__class__.__name__ in excluded:
modules/shared.py
+2 -3
View File
@@ -18,8 +18,8 @@ from modules.paths import models_path, script_path, data_path, sd_configs_path,
from modules.dml import memory_providers, default_memory_provider, directml_do_hijack
from modules.onnx_impl import initialize_onnx, execution_providers
from modules.memstats import memory_stats, ram_stats # pylint: disable=unused-import
from modules.interrogate.legacy import category_types
from modules.ui_components import DropdownEditable
import modules.interrogate
import modules.memmon
import modules.styles
import modules.paths as paths
@@ -36,7 +36,6 @@ cmd_opts, _ = parser.parse_known_args()
hide_dirs = {"visible": not cmd_opts.hide_ui_dir_config}
xformers_available = False
locking_available = True # used by file read/write locking
interrogator = modules.interrogate.InterrogateModels(os.path.join("models", "interrogate"))
sd_upscalers = []
detailers = []
face_restorers = []
@@ -896,7 +895,7 @@ options_templates.update(options_section(('interrogate', "Interrogate"), {
"interrogate_clip_num_beams": OptionInfo(1, "Interrogate: num_beams for BLIP", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1}),
"interrogate_clip_min_length": OptionInfo(32, "Interrogate: minimum description length", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1}),
"interrogate_clip_max_length": OptionInfo(192, "Interrogate: maximum description length", gr.Slider, {"minimum": 1, "maximum": 256, "step": 1}),
"interrogate_clip_skip_categories": OptionInfo(["artists", "movements", "flavors"], "Interrogate: skip categories", gr.CheckboxGroup, lambda: {"choices": modules.interrogate.category_types()}, refresh=modules.interrogate.category_types),
"interrogate_clip_skip_categories": OptionInfo(["artists", "movements", "flavors"], "Interrogate: skip categories", gr.CheckboxGroup, lambda: {"choices": category_types()}, refresh=category_types),
"interrogate_deepbooru_score_threshold": OptionInfo(0.65, "Interrogate: deepbooru score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
"deepbooru_sort_alpha": OptionInfo(False, "Interrogate: deepbooru sort alphabetically"),
"deepbooru_use_spaces": OptionInfo(False, "Use spaces for tags in deepbooru"),
modules/ui_common.py
+6 -5
View File
@@ -221,19 +221,20 @@ def open_folder(result_gallery, gallery_index = 0):
subprocess.Popen(["xdg-open", path]) # pylint: disable=consider-using-with
def interrogate_clip(image):
def interrogate_clip(image): # legacy function
if image is None:
shared.log.error("Interrogate: no image selected")
return gr.update()
prompt = shared.interrogator.interrogate(image)
from modules.interrogate import legacy
prompt = legacy.interrogator.interrogate(image)
return gr.update() if prompt is None else prompt
def interrogate_booru(image):
def interrogate_booru(image): # legacy function
if image is None:
shared.log.error("Interrogate: no image selected")
return gr.update()
from modules import deepbooru
from modules.interrogate import deepbooru
prompt = deepbooru.model.tag(image)
return gr.update() if prompt is None else prompt
@@ -258,7 +259,7 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None):
elem_classes=["gallery_main"],
)
if prompt is not None:
ui_sections.create_interrogate_button(tab=tabname, inputs=result_gallery, output=prompt)
ui_sections.create_interrogate_button(tab=tabname, inputs=result_gallery, outputs=prompt)
# interrogate_clip_btn, interrogate_booru_btn = ui_sections.create_interrogate_buttons(tabname)
# interrogate_clip_btn.click(fn=interrogate_clip, inputs=[result_gallery], outputs=[prompt])
# interrogate_booru_btn.click(fn=interrogate_booru, inputs=[result_gallery], outputs=[prompt])
modules/ui_control_helpers.py
+5 -4
View File
@@ -47,19 +47,20 @@ def initialize():
scripts.scripts_control.initialize_scripts(is_img2img=False, is_control=True)
def interrogate_clip():
def interrogate_clip(): # legacy function
prompt = None
try:
prompt = shared.interrogator.interrogate(input_source[0])
from modules.interrogate import legacy
prompt = legacy.interrogator.interrogate(input_source[0])
except Exception:
pass
return gr.update() if prompt is None else prompt
def interrogate_booru():
def interrogate_booru(): # legacy function
prompt = None
try:
from modules import deepbooru
from modules.interrogate import deepbooru
prompt = deepbooru.model.tag(input_source[0])
except Exception:
pass
modules/ui_postprocessing.py
+15 -14
View File
@@ -1,7 +1,8 @@
import json
import gradio as gr
from modules import scripts, shared, ui_common, postprocessing, call_queue, interrogate, generation_parameters_copypaste
from modules import scripts, shared, ui_common, postprocessing, call_queue, generation_parameters_copypaste
from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call, wrap_gradio_call # pylint: disable=unused-import
from modules.interrogate import legacy
def submit_info(image):
@@ -46,8 +47,8 @@ def create_ui():
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')
ui_common.create_refresh_button(clip_model, interrogate.get_clip_models, lambda: {"choices": interrogate.get_clip_models()}, 'refresh_interrogate_models')
blip_model = gr.Dropdown(list(interrogate.caption_models), value='blip-base', label='Caption model')
ui_common.create_refresh_button(clip_model, legacy.get_clip_models, lambda: {"choices": legacy.get_clip_models()}, 'refresh_interrogate_models')
blip_model = gr.Dropdown(list(legacy.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():
@@ -56,20 +57,20 @@ def create_ui():
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=[])
caption_max_length.change(fn=legacy.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
chunk_size.change(fn=legacy.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
min_flavors.change(fn=legacy.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
max_flavors.change(fn=legacy.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[])
flavor_intermediate_count.change(fn=legacy.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, 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)
btn_interrogate_img.click(legacy.interrogate_image, inputs=[image, clip_model, blip_model, mode], outputs=prompt)
btn_analyze_img.click(legacy.analyze_image, inputs=[image, clip_model, blip_model], outputs=[medium, artist, movement, trending, flavor])
btn_unload.click(legacy.unload_clip_model)
with gr.Tab("Interrogate 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)
@@ -81,11 +82,11 @@ def create_ui():
batch = gr.Text(label="Prompts", lines=10)
with gr.Row():
clip_model = gr.Dropdown([], value='ViT-L-14/openai', label='CLiP Batch Model')
ui_common.create_refresh_button(clip_model, interrogate.get_clip_models, lambda: {"choices": interrogate.get_clip_models()}, 'refresh_interrogate_models')
ui_common.create_refresh_button(clip_model, legacy.get_clip_models, lambda: {"choices": legacy.get_clip_models()}, 'refresh_interrogate_models')
with gr.Row(elem_id='interrogate_buttons_batch'):
btn_interrogate_batch = gr.Button("Interrogate", elem_id="interrogate_btn_interrogate", variant='primary')
with gr.Tab("Visual Query"):
from modules import vqa
from modules.interrogate import vqa
with gr.Row():
vqa_image = gr.Image(type='pil', label="Image")
with gr.Row():
@@ -148,7 +149,7 @@ def create_ui():
]
)
btn_interrogate_batch.click(
fn=interrogate.interrogate_batch,
fn=legacy.interrogate_batch,
inputs=[batch_files, batch_folder, batch_str, clip_model, blip_model, mode, save_output],
outputs=[batch],
)
modules/ui_sections.py
+3 -3
View File
@@ -1,6 +1,7 @@
import gradio as gr
from modules import shared, modelloader, ui_symbols, ui_common, sd_samplers
from modules.ui_components import ToolButton
from modules.interrogate import interrogate
def create_toprow(is_img2img: bool = False, id_part: str = None):
@@ -92,11 +93,10 @@ def create_resolution_inputs(tab):
def create_interrogate_button(tab: str, inputs: list, outputs: str):
button_interrogate = gr.Button(ui_symbols.interrogate, elem_id=f"{tab}_interrogate", elem_classes=['interrogate'])
button_interrogate.click(fn=lambda: None, _js='() => quickInterrogate()', inputs=[], outputs=[])
return button_interrogate
button_interrogate.click(fn=interrogate.interrogate, inputs=inputs, outputs=[outputs])
def create_interrogate_buttons(tab):
def create_interrogate_buttons(tab): # legacy function
button_interrogate = gr.Button(ui_symbols.int_clip, elem_id=f"{tab}_interrogate", elem_classes=['interrogate-clip'])
button_deepbooru = gr.Button(ui_symbols.int_blip, elem_id=f"{tab}_deepbooru", elem_classes=['interrogate-blip'])
return button_interrogate, button_deepbooru
scripts/loopback.py
+4 -2
View File
@@ -2,7 +2,7 @@ import math
import gradio as gr
import modules.scripts as scripts
from modules import deepbooru, images, processing, shared
from modules import images, processing
from modules.processing import Processed
from modules.shared import opts, state
@@ -90,8 +90,10 @@ class Script(scripts.Script):
if append_interrogation != "None":
p.prompt = f"{original_prompt}, " if original_prompt else ""
if append_interrogation == "CLIP":
p.prompt += shared.interrogator.interrogate(p.init_images[0])
from modules.interrogate import legacy
p.prompt += legacy.interrogator.interrogate(p.init_images[0])
elif append_interrogation == "DeepBooru":
from modules.interrogate import deepbooru
p.prompt += deepbooru.model.tag(p.init_images[0])
state.job = f"loopback iteration {i+1}/{loops} batch {n+1}/{batch_count}"