diff --git a/CHANGELOG.md b/CHANGELOG.md index 4957fb1c2..d3d64904d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,11 @@ - add support for stream-loading, this can speed up model loading when models are located on network drives *set in settings -> models & loading -> model load using streams* - enhanced error logging +- **Interrogate/Captioning** + - single interrogate button for every input or output image + - behavior of interrogate configurable in *settings -> interrogate* + with detailed defaults for each model type also configurable + - select between 100+ *OpenCLiP* supported models, 10+ built-in *VLMs*, *DeepDanbooru* - **Other**: - **Networks**: imporove search/filter and add visual indicators for types - **balanced offload** new defaults: *lowvram/4gb min threshold: 0, medvram/8gb min threshold: 0, default min threshold 0.25* @@ -45,7 +50,8 @@ - ipex device wrapper with adetailer - openvino error handling - relax python version checks for rocm - - simplify and improve file wildcard matching + - simplify and improve file wildcard matching + - fix `rich` version ## Update for 2025-01-29 diff --git a/html/locale_en.json b/html/locale_en.json index 0767ccd23..8c7f036b4 100644 --- a/html/locale_en.json +++ b/html/locale_en.json @@ -23,6 +23,7 @@ {"id":"","label":"🔍","localized":"","hint":"Search"}, {"id":"","label":"🖌️","localized":"","hint":"LaMa remove selected object from image"}, {"id":"","label":"🖼️","localized":"","hint":"Show preview"}, + {"id":"","label":"♻","localized":"","hint":"Interrogate image"}, {"id":"","label":"✎","localized":"","hint":"Interrogate image using BLIP model"}, {"id":"","label":"✐","localized":"","hint":"Interrogate image using DeepBooru model"}, {"id":"","label":"↶","localized":"","hint":"Apply selected style to prompt"}, diff --git a/installer.py b/installer.py index fed40d835..0e7c4c16a 100644 --- a/installer.py +++ b/installer.py @@ -559,11 +559,10 @@ def install_cuda(): log.info('CUDA: nVidia toolkit detected') ts('cuda', t_start) if args.use_nightly: - cmd = os.environ.get('TORCH_COMMAND', '--pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu126') + cmd = os.environ.get('TORCH_COMMAND', 'pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 --extra-index-url https://download.pytorch.org/whl/nightly/cu126') else: # cmd = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cu124 torchvision==0.20.1+cu124 --index-url https://download.pytorch.org/whl/cu124') cmd = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+cu126 torchvision==0.21.0+cu126 --index-url https://download.pytorch.org/whl/cu126') - # TODO blackwell requires cuda==12.8 torch release is pending return cmd diff --git a/javascript/sdnext.css b/javascript/sdnext.css index b6f72b593..42686c88a 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -111,9 +111,10 @@ button.custom-button { border-radius: var(--button-large-radius); padding: var(- #txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt, #control_prompt, #control_neg_prompt { display: contents; } #txt2img_actions_column, #img2img_actions_column, #control_actions { flex-flow: wrap; justify-content: space-between; } +.interrogate { position: absolute; right: 2.8em; top: 0.2em; max-width: fit-content; background: none !important; z-index: 50; font-size: 1.5em !important; } +.interrogate:hover { background: var(--button-primary-background-fill-hover) !important; } .interrogate-clip { position: absolute; right: 6em; top: 8px; max-width: fit-content; background: none !important; z-index: 50; } .interrogate-blip { position: absolute; right: 4em; top: 8px; max-width: fit-content; background: none !important; z-index: 50; } -.interrogate { position: absolute; right: 4em; top: 8px; max-width: fit-content; background: none !important; z-index: 50; } .interrogate-col { min-width: 0 !important; max-width: fit-content; margin-right: var(--spacing-xxl); } .interrogate-col>button { flex: 1; width: 7em; max-height: 84px; } #sampler_selection_img2img { margin-top: 1em; } diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py index 8112fe4a8..52c4b5301 100644 --- a/modules/api/endpoints.py +++ b/modules/api/endpoints.py @@ -74,8 +74,8 @@ def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, f return res def get_interrogate(): - from modules.interrogate.legacy import get_clip_models - return ['clip', 'deepdanbooru'] + get_clip_models() + from modules.interrogate.openclip import refresh_clip_models + return ['clip', 'deepdanbooru'] + refresh_clip_models() def post_interrogate(req: models.ReqInterrogate): if req.image is None or len(req.image) < 64: @@ -84,8 +84,8 @@ def post_interrogate(req: models.ReqInterrogate): image = image.convert('RGB') if req.model == "clip": try: - from modules.interrogate import legacy - caption = legacy.interrogator.interrogate(image) + from modules.interrogate import openclip + caption = openclip.interrogator.interrogate(image) except Exception as e: caption = str(e) return models.ResInterrogate(caption=caption) @@ -94,8 +94,8 @@ def post_interrogate(req: models.ReqInterrogate): caption = deepbooru.model.tag(image) return models.ResInterrogate(caption=caption) else: - from modules.interrogate.legacy import interrogate_image, analyze_image, get_clip_models - if req.model not in get_clip_models(): + from modules.interrogate.openclip import interrogate_image, analyze_image, refresh_clip_models + if req.model not in refresh_clip_models(): raise HTTPException(status_code=404, detail="Model not found") try: caption = interrogate_image(image, clip_model=req.clip_model, blip_model=req.blip_model, mode=req.mode) diff --git a/modules/interrogate/deepbooru.py b/modules/interrogate/deepbooru.py index 5e54dcc1a..24d7aec51 100644 --- a/modules/interrogate/deepbooru.py +++ b/modules/interrogate/deepbooru.py @@ -4,7 +4,7 @@ import threading import torch import numpy as np from PIL import Image -from modules import modelloader, paths, devices, images, shared +from modules import modelloader, paths, devices, shared re_special = re.compile(r'([\\()])') load_lock = threading.Lock() @@ -19,7 +19,7 @@ class DeepDanbooru: if self.model is not None: return model_path = os.path.join(paths.models_path, "DeepDanbooru") - shared.log.debug(f'Load interrogate model: type=DeepDanbooru folder="{model_path}"') + shared.log.debug(f'Interrogate load: module=DeepDanbooru folder="{model_path}"') files = modelloader.load_models( model_path=model_path, model_url='https://github.com/AUTOMATIC1111/TorchDeepDanbooru/releases/download/v1/model-resnet_custom_v3.pt', @@ -30,7 +30,6 @@ class DeepDanbooru: 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() self.model.to(devices.cpu, devices.dtype) @@ -39,9 +38,9 @@ class DeepDanbooru: self.model.to(devices.device) def stop(self): - if not shared.opts.interrogate_keep_models_in_memory: + if shared.opts.interrogate_offload: self.model.to(devices.cpu) - devices.torch_gc() + devices.torch_gc() def tag(self, pil_image): self.start() @@ -57,14 +56,14 @@ class DeepDanbooru: pil_image = Image.open(pil_image['name']) if pil_image is None: return '' - pic = images.resize_image(2, pil_image.convert("RGB"), 512, 512) + pic = pil_image.resize((512, 512), resample=Image.Resampling.LANCZOS).convert("RGB") a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255 with devices.inference_context(), devices.autocast(): x = torch.from_numpy(a).to(devices.device) y = self.model(x)[0].detach().float().cpu().numpy() probability_dict = {} for tag, probability in zip(self.model.tags, y): - if probability < shared.opts.interrogate_deepbooru_score_threshold: + if probability < shared.opts.deepbooru_score_threshold: continue if tag.startswith("rating:"): continue @@ -82,9 +81,11 @@ class DeepDanbooru: tag_outformat = tag_outformat.replace('_', ' ') if shared.opts.deepbooru_escape: tag_outformat = re.sub(re_special, r'\\\1', tag_outformat) - if shared.opts.interrogate_return_ranks and not force_disable_ranks: + if shared.opts.deepbooru_clip_score and not force_disable_ranks: tag_outformat = f"({tag_outformat}:{probability:.3f})" res.append(tag_outformat) + if len(res) > shared.opts.deepbooru_max_tags: + res = res[:shared.opts.deepbooru_max_tags] return ", ".join(res) diff --git a/modules/interrogate/interrogate.py b/modules/interrogate/interrogate.py index 1696e88e7..45d485cef 100644 --- a/modules/interrogate/interrogate.py +++ b/modules/interrogate/interrogate.py @@ -1,4 +1,36 @@ +import time +from PIL import Image +from modules import shared + + def interrogate(image): - from modules.interrogate import legacy - prompt = legacy.interrogator.interrogate(image) - return prompt + if isinstance(image, list): + image = image[0] if len(image) > 0 else None + if isinstance(image, dict) and 'name' in image: + image = Image.open(image['name']) + if image is None: + return '' + t0 = time.time() + if shared.opts.interrogate_default_type == 'OpenCLiP': + shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} clip="{shared.opts.interrogate_clip_model}" blip="{shared.opts.interrogate_blip_model}" mode="{shared.opts.interrogate_clip_mode}"') + from modules.interrogate import openclip + openclip.load_interrogator(clip_model=shared.opts.interrogate_clip_model, blip_model=shared.opts.interrogate_blip_model) + openclip.update_interrogate_params() + prompt = openclip.interrogate(image, mode=shared.opts.interrogate_clip_mode) + shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"') + return prompt + elif shared.opts.interrogate_default_type == 'DeepBooru': + shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type}') + from modules.interrogate import deepbooru + prompt = deepbooru.model.tag(image) + shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"') + return prompt + elif shared.opts.interrogate_default_type == 'VLM': + shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} vlm="{shared.opts.interrogate_vlm_model}" prompt="{shared.opts.interrogate_vlm_prompt}"') + from modules.interrogate import vqa + prompt = vqa.interrogate(image=image, model_name=shared.opts.interrogate_vlm_model, question=shared.opts.interrogate_vlm_prompt) + shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"') + return prompt + else: + shared.log.error(f'Interrogate: type="{shared.opts.interrogate_default_type}" unknown') + return '' diff --git a/modules/interrogate/legacy.py b/modules/interrogate/openclip.py similarity index 84% rename from modules/interrogate/legacy.py rename to modules/interrogate/openclip.py index 486a6e2bb..3b716571a 100644 --- a/modules/interrogate/legacy.py +++ b/modules/interrogate/openclip.py @@ -1,6 +1,5 @@ import os import sys -import time from collections import namedtuple from pathlib import Path import threading @@ -14,7 +13,7 @@ from modules import devices, paths, shared, lowvram, errors, sd_models config = { - "caption_max_length": 64, + "caption_max_length": 74, "chunk_size": 1024, "flavor_intermediate_count": 1024, "min_flavors": 2, @@ -30,6 +29,14 @@ caption_models = { 'blip2-flip-t5-xl': 'Salesforce/blip2-flan-t5-xl', 'blip2-flip-t5-xxl': 'Salesforce/blip2-flan-t5-xxl', } +caption_types = [ + 'best', + 'fast', + 'classic', + 'caption', + 'negative', +] +clip_models = [] ci = None blip_image_eval_size = 384 clip_model_name = 'ViT-L/14' @@ -43,7 +50,7 @@ def category_types(): def download_default_clip_interrogate_categories(content_dir): - shared.log.info("Downloading CLIP categories...") + shared.log.info("Interrogate: downloading CLIP categories...") tmpdir = f"{content_dir}_tmp" cat_types = ["artists", "flavors", "mediums", "movements"] try: @@ -106,7 +113,7 @@ class InterrogateModels: import modules.modelloader as modelloader model_path = os.path.join(paths.models_path, "BLIP") download_name='model_base_caption_capfilt_large.pth' - shared.log.debug(f'Model interrogate load: type=BLiP model={download_name} path={model_path}') + shared.log.debug(f'Interrogate load: module=BLiP model="{download_name}" path="{model_path}"') files = modelloader.load_models( model_path=model_path, model_url='https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_caption_capfilt_large.pth', @@ -119,7 +126,7 @@ class InterrogateModels: def load_clip_model(self): with load_lock: - shared.log.debug(f'Model interrogate load: type=CLiP model={clip_model_name} path={shared.opts.clip_models_path}') + shared.log.debug(f'Interrogate load: module=CLiP model="{clip_model_name}" path="{shared.opts.clip_models_path}"') import clip if self.running_on_cpu: model, preprocess = clip.load(clip_model_name, device="cpu", download_root=shared.opts.clip_models_path) @@ -143,12 +150,12 @@ class InterrogateModels: self.dtype = next(self.clip_model.parameters()).dtype def send_clip_to_ram(self): - if not shared.opts.interrogate_keep_models_in_memory: + if shared.opts.interrogate_offload: if self.clip_model is not None: self.clip_model = self.clip_model.to(devices.cpu) def send_blip_to_ram(self): - if not shared.opts.interrogate_keep_models_in_memory: + if shared.opts.interrogate_offload: if self.blip_model is not None: self.blip_model = self.blip_model.to(devices.cpu) @@ -180,36 +187,33 @@ class InterrogateModels: transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)) ])(pil_image).unsqueeze(0).type(self.dtype).to(devices.device) with devices.inference_context(): - caption = self.blip_model.generate(gpu_image, sample=False, num_beams=shared.opts.interrogate_clip_num_beams, min_length=shared.opts.interrogate_clip_min_length, max_length=shared.opts.interrogate_clip_max_length) + min_length = min(shared.opts.interrogate_clip_min_length, shared.opts.interrogate_clip_max_length) + max_length = max(shared.opts.interrogate_clip_min_length, shared.opts.interrogate_clip_max_length) + caption = self.blip_model.generate(gpu_image, sample=False, num_beams=shared.opts.interrogate_clip_num_beams, min_length=min_length, max_length=max_length) return caption[0] - def interrogate(self, pil_image): + def interrogate(self, image): res = "" 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() self.load() - if isinstance(pil_image, list): - pil_image = pil_image[0] if len(pil_image) > 0 else None - if isinstance(pil_image, dict) and 'name' in pil_image: - pil_image = Image.open(pil_image['name']) - if pil_image is None: + if isinstance(image, list): + image = image[0] if len(image) > 0 else None + if isinstance(image, dict) and 'name' in image: + image = Image.open(image['name']) + if image is None: return '' - pil_image = pil_image.convert("RGB") - caption = self.generate_caption(pil_image) - self.send_blip_to_ram() - devices.torch_gc() + image = image.convert("RGB") + caption = self.generate_caption(image) res = caption - clip_image = self.clip_preprocess(pil_image).unsqueeze(0).type(self.dtype).to(devices.device) + clip_image = self.clip_preprocess(image).unsqueeze(0).type(self.dtype).to(devices.device) with devices.inference_context(), devices.autocast(): image_features = self.clip_model.encode_image(clip_image).type(self.dtype) image_features /= image_features.norm(dim=-1, keepdim=True) for _name, topn, items in self.categories(): matches = self.rank(image_features, items, top_count=topn) for match, score in matches: - if shared.opts.interrogate_return_ranks: + if shared.opts.interrogate_clip_score: res += f", ({match}:{score/100:.3f})" else: res += f", {match}" @@ -237,23 +241,34 @@ 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) +def update_interrogate_params(caption_max_length:int=None, chunk_size:int=None, min_flavors:int=None, max_flavors:int=None, flavor_intermediate_count:int=None): + config["caption_max_length"] = int(caption_max_length or shared.opts.interrogate_clip_max_length) + config["clip_offload"] = shared.opts.interrogate_offload + config["caption_offload"] = shared.opts.interrogate_offload + config["min_flavors"] = int(min_flavors or shared.opts.interrogate_clip_min_flavors) + config["max_flavors"] = int(max_flavors or shared.opts.interrogate_clip_max_flavors) + if chunk_size is not None: + config["chunk_size"] = int(chunk_size) + if flavor_intermediate_count is not None: + 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}') + shared.log.debug(f'Interrogate: type={shared.opts.interrogate_default_type} config={config}') + def get_clip_models(): + return clip_models + + +def refresh_clip_models(): + global clip_models # pylint: disable=global-statement import open_clip models = sorted(open_clip.list_pretrained()) - shared.log.info(f'Interrogate: pkg=openclip version={open_clip.__version__} models={len(models)}') - return ['/'.join(x) for x in models] + shared.log.debug(f'Interrogate: pkg=openclip version={open_clip.__version__} models={len(models)}') + clip_models = ['/'.join(x) for x in models] + return clip_models def load_interrogator(clip_model, blip_model): @@ -263,6 +278,7 @@ def load_interrogator(clip_model, blip_model): clip_interrogator.clip_interrogator.CAPTION_MODELS = caption_models global ci # pylint: disable=global-statement if ci is None: + shared.log.debug(f'Interrogate load: clip="{clip_model}" blip="{blip_model}"') interrogator_config = clip_interrogator.Config( device=devices.get_optimal_device(), cache_path=os.path.join(paths.models_path, 'Interrogator'), @@ -275,25 +291,18 @@ def load_interrogator(clip_model, blip_model): 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(): - if ci is not None: - shared.log.debug('Interrogate offload') + if ci is not None and shared.opts.interrogate_offload: ci.caption_model = ci.caption_model.to(devices.cpu) ci.clip_model = ci.clip_model.to(devices.cpu) ci.caption_offloaded = True @@ -302,8 +311,13 @@ def unload_clip_model(): def interrogate(image, mode, caption=None): - shared.log.info(f'Interrogate: mode={mode} image={image}') - t0 = time.time() + if isinstance(image, list): + image = image[0] if len(image) > 0 else None + if isinstance(image, dict) and 'name' in image: + image = Image.open(image['name']) + if image is None: + return '' + image = image.convert("RGB") if mode == 'best': prompt = ci.interrogate(image, caption=caption, min_flavors=config["min_flavors"], max_flavors=config["max_flavors"]) elif mode == 'caption': @@ -316,8 +330,6 @@ def interrogate(image, mode, caption=None): 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 diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index d0172159b..c9d570681 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -9,7 +9,7 @@ from modules import shared, devices, errors processor = None model = None loaded: str = None -MODELS = { +vlm_models = { "MS Florence 2 Base": "microsoft/Florence-2-base", # 0.5GB "MS Florence 2 Large": "microsoft/Florence-2-large", # 1.5GB "MiaoshouAI PromptGen 1.5 Base": "MiaoshouAI/Florence-2-base-PromptGen-v1.5@c06a5f02cc6071a5d65ee5d294cf3732d3097540", # 1.1GB @@ -27,17 +27,31 @@ MODELS = { "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB "Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB } +vlm_prompts = [ + '', + '', + '', + '', + '', + '', + '', + '', + '', + '', + '', + '', + '', +] def git(question: str, image: Image.Image, repo: str = None): global processor, model, loaded # pylint: disable=global-statement if model is None or loaded != repo: + shared.log.debug(f'Interrogate load: vlm="{repo}"') model = transformers.GitForCausalLM.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) processor = transformers.GitProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) loaded = repo model.to(devices.device, devices.dtype) - shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}') - pixel_values = processor(images=image, return_tensors="pt").pixel_values git_dict = {} git_dict['pixel_values'] = pixel_values.to(devices.device, devices.dtype) @@ -49,14 +63,13 @@ def git(question: str, image: Image.Image, repo: str = None): with devices.inference_context(): generated_ids = model.generate(**git_dict) response = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] - - shared.log.debug(f'VQA: response={response}') return response def blip(question: str, image: Image.Image, repo: str = None): global processor, model, loaded # pylint: disable=global-statement if model is None or loaded != repo: + shared.log.debug(f'Interrogate load: vlm="{repo}"') model = transformers.BlipForQuestionAnswering.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) processor = transformers.BlipProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) loaded = repo @@ -66,21 +79,17 @@ def blip(question: str, image: Image.Image, repo: str = None): with devices.inference_context(): outputs = model.generate(**inputs) response = processor.decode(outputs[0], skip_special_tokens=True) - - model.to(devices.cpu) - shared.log.debug(f'VQA: response={response}') return response def vilt(question: str, image: Image.Image, repo: str = None): global processor, model, loaded # pylint: disable=global-statement if model is None or loaded != repo: + shared.log.debug(f'Interrogate load: vlm="{repo}"') model = transformers.ViltForQuestionAnswering.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) processor = transformers.ViltProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) loaded = repo model.to(devices.device) - shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}') - inputs = processor(image, question, return_tensors="pt") inputs = inputs.to(devices.device) with devices.inference_context(): @@ -88,20 +97,17 @@ def vilt(question: str, image: Image.Image, repo: str = None): logits = outputs.logits idx = logits.argmax(-1).item() response = model.config.id2label[idx] - - shared.log.debug(f'VQA: response={response}') return response def pix(question: str, image: Image.Image, repo: str = None): global processor, model, loaded # pylint: disable=global-statement if model is None or loaded != repo: + shared.log.debug(f'Interrogate load: vlm="{repo}"') model = transformers.Pix2StructForConditionalGeneration.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) processor = transformers.Pix2StructProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) loaded = repo model.to(devices.device) - shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}') - if len(question) > 0: inputs = processor(images=image, text=question, return_tensors="pt").to(devices.device) else: @@ -109,14 +115,13 @@ def pix(question: str, image: Image.Image, repo: str = None): with devices.inference_context(): outputs = model.generate(**inputs) response = processor.decode(outputs[0], skip_special_tokens=True) - - shared.log.debug(f'VQA: response={response}') return response def moondream(question: str, image: Image.Image, repo: str = None): global processor, model, loaded # pylint: disable=global-statement if model is None or loaded != repo: + shared.log.debug(f'Interrogate load: vlm="{repo}"') model = transformers.AutoModelForCausalLM.from_pretrained( repo, revision="2024-08-26", @@ -127,15 +132,12 @@ def moondream(question: str, image: Image.Image, repo: str = None): loaded = repo model.eval() model.to(devices.device, devices.dtype) - shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}') - if len(question) < 2: question = "Describe the image." + question = question.replace('<', '').replace('>', '') encoded = model.encode_image(image) with devices.inference_context(): response = model.answer_question(encoded, question, processor) - - shared.log.debug(f'VQA: response="{response}"') return response @@ -148,6 +150,7 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str R.remove("flash_attn") # flash_attn is optional return R if model is None or loaded != repo: + shared.log.debug(f'Interrogate load: vlm="{repo}"') transformers.dynamic_module_utils.get_imports = get_imports model = transformers.AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, revision=revision, cache_dir=shared.opts.hfcache_dir) processor = transformers.AutoProcessor.from_pretrained(repo, trust_remote_code=True, revision=revision, cache_dir=shared.opts.hfcache_dir) @@ -155,8 +158,6 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str loaded = repo model.eval() model.to(devices.device, devices.dtype) - shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}') - if question.startswith('<'): task = question.split('>', 1)[0] + '>' else: @@ -169,13 +170,12 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str generated_ids = model.generate( input_ids=input_ids, pixel_values=pixel_values, - max_new_tokens=1024, - num_beams=3, + max_new_tokens=shared.opts.interrogate_vlm_max_length, + num_beams=shared.opts.interrogate_vlm_num_beams, do_sample=False ) generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] response = processor.post_process_generation(generated_text, task="task", image_size=(image.width, image.height)) - if 'task' in response: response = response['task'] if 'answer' in response: @@ -183,44 +183,48 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str if isinstance(response, dict): response = json.dumps(response) response = response.replace('\n', '').replace('\r', '').replace('\t', '').strip() - shared.log.debug(f'VQA: task={task} response="{response}"') return response -def interrogate(vqa_question, vqa_image, vqa_model_req): +def interrogate(question, image, model_name): + if isinstance(image, list): + image = image[0] if len(image) > 0 else None + if isinstance(image, dict) and 'name' in image: + image = Image.open(image['name']) + if image is None: + return '' try: - vqa_model = MODELS.get(vqa_model_req, None) + vqa_model = vlm_models.get(model_name, None) revision = None if '@' in vqa_model: vqa_model, revision = vqa_model.split('@') - shared.log.debug(f'VQA: model="{vqa_model}" question="{vqa_question}" image={vqa_image}') - if vqa_image is None: + if image is None: answer = 'no image provided' return answer - if vqa_model_req is None: + if model_name is None: answer = 'no model selected' return answer if vqa_model is None: - answer = f'unknown: model={vqa_model_req} available={MODELS.keys()}' + answer = f'unknown: model={model_name} available={vlm_models.keys()}' return answer if 'git' in vqa_model.lower(): - answer = git(vqa_question, vqa_image, vqa_model) + answer = git(question, image, vqa_model) elif 'vilt' in vqa_model.lower(): - answer = vilt(vqa_question, vqa_image, vqa_model) + answer = vilt(question, image, vqa_model) elif 'blip' in vqa_model.lower(): - answer = blip(vqa_question, vqa_image, vqa_model) + answer = blip(question, image, vqa_model) elif 'pix' in vqa_model.lower(): - answer = pix(vqa_question, vqa_image, vqa_model) + answer = pix(question, image, vqa_model) elif 'moondream2' in vqa_model.lower(): - answer = moondream(vqa_question, vqa_image, vqa_model) + answer = moondream(question, image, vqa_model) elif 'florence' in vqa_model.lower(): - answer = florence(vqa_question, vqa_image, vqa_model, revision) + answer = florence(question, image, vqa_model, revision) else: answer = 'unknown model' except Exception as e: errors.display(e, 'VQA') answer = 'error' - if model is not None: + if shared.opts.interrogate_offload and model is not None: model.to(devices.cpu) devices.torch_gc() return answer diff --git a/modules/shared.py b/modules/shared.py index 4aa5a87ce..9f0fc4aa9 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -18,7 +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.interrogate.openclip import caption_models, caption_types, get_clip_models, refresh_clip_models, category_types +from modules.interrogate.vqa import vlm_models, vlm_prompts from modules.ui_components import DropdownEditable import modules.memmon import modules.styles @@ -890,17 +891,35 @@ options_templates.update(options_section(('control', "Control Options"), { })) options_templates.update(options_section(('interrogate', "Interrogate"), { - "interrogate_keep_models_in_memory": OptionInfo(False, "Interrogate: keep models in VRAM"), - "interrogate_return_ranks": OptionInfo(True, "Interrogate: include ranks of model tags matches in results"), - "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": 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"), - "deepbooru_escape": OptionInfo(True, "Escape brackets in deepbooru"), - "deepbooru_filter_tags": OptionInfo("", "Filter out tags from deepbooru output"), + "interrogate_default_type": OptionInfo("OpenCLiP", "Default type", gr.Radio, {"choices": ["OpenCLiP", "VLM", "DeepBooru"]}), + "interrogate_offload": OptionInfo(True, "Interrogate: offload models "), + + "interrogate_clip_sep": OptionInfo("

OpenCLiP

", "", gr.HTML), + "interrogate_clip_model": OptionInfo("ViT-L-14/openai", "CLiP: default model", gr.Dropdown, lambda: {"choices": get_clip_models()}, refresh=refresh_clip_models), + "interrogate_clip_mode": OptionInfo(caption_types[0], "CLiP: default mode", gr.Dropdown, {"choices": caption_types}), + "interrogate_blip_model": OptionInfo(list(caption_models)[0], "CLiP: default captioner", gr.Dropdown, {"choices": list(caption_models)}), + "interrogate_clip_score": OptionInfo(False, "CLiP: include scores in results"), + "interrogate_clip_num_beams": OptionInfo(1, "CLiP: num beams", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1}), + "interrogate_clip_min_length": OptionInfo(32, "CLiP: min length", gr.Slider, {"minimum": 1, "maximum": 128, "step": 1}), + "interrogate_clip_max_length": OptionInfo(74, "CLiP: max length", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1}), + "interrogate_clip_min_flavors": OptionInfo(2, "CLiP: min flavors", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}), + "interrogate_clip_max_flavors": OptionInfo(8, "CLiP: max flavors", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}), + "interrogate_clip_skip_categories": OptionInfo(["artists", "movements", "flavors"], "CLiP: skip categories", gr.CheckboxGroup, lambda: {"choices": category_types()}, refresh=category_types), + + "interrogate_vlm_sep": OptionInfo("

VLM

", "", gr.HTML), + "interrogate_vlm_model": OptionInfo(list(vlm_models)[0], "VLM: default model", gr.Dropdown, {"choices": list(vlm_models)}), + "interrogate_vlm_prompt": OptionInfo(vlm_prompts[2], "VLM: default prompt", DropdownEditable, {"choices": vlm_prompts }), + "interrogate_vlm_num_beams": OptionInfo(3, "VLM: num beams", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1}), + "interrogate_vlm_max_length": OptionInfo(512, "VLM: max length", gr.Slider, {"minimum": 1, "maximum": 4096, "step": 1}), + + "deepbooru_sep": OptionInfo("

DeepBooru

", "", gr.HTML), + "deepbooru_score_threshold": OptionInfo(0.65, "DeepBooru: score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}), + "deepbooru_max_tags": OptionInfo(74, "DeepBooru: max tags", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1}), + "deepbooru_clip_score": OptionInfo(False, "DeepBooru: include scores in results"), + "deepbooru_sort_alpha": OptionInfo(False, "DeepBooru: sort alphabetically"), + "deepbooru_use_spaces": OptionInfo(False, "DeepBooru: use spaces for tags"), + "deepbooru_escape": OptionInfo(True, "DeepBooru: escape brackets"), + "deepbooru_filter_tags": OptionInfo("", "DeepBooru: exclude tags"), })) options_templates.update(options_section(('huggingface', "Huggingface"), { diff --git a/modules/ui_common.py b/modules/ui_common.py index c3d5f80fd..efc40e1e0 100644 --- a/modules/ui_common.py +++ b/modules/ui_common.py @@ -225,8 +225,8 @@ def interrogate_clip(image): # legacy function if image is None: shared.log.error("Interrogate: no image selected") return gr.update() - from modules.interrogate import legacy - prompt = legacy.interrogator.interrogate(image) + from modules.interrogate import openclip + prompt = openclip.interrogator.interrogate(image) return gr.update() if prompt is None else prompt @@ -260,10 +260,6 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None): ) if prompt is not None: 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]) - with gr.Column(elem_id=f"{tabname}_footer", elem_classes="gallery_footer"): dummy_component = gr.Label(visible=False) diff --git a/modules/ui_control.py b/modules/ui_control.py index 53a44f0e0..6f92c6cda 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -191,7 +191,7 @@ def create_ui(_blocks: gr.Blocks=None): input_image = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="editor", height=gr_height, visible=True, image_mode='RGB', elem_id='control_input_select', elem_classes=['control-image']) input_resize = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="select", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_resize', elem_classes=['control-image']) input_inpaint = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="sketch", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_inpaint', brush_radius=32, mask_opacity=0.6, elem_classes=['control-image']) - btn_interrogate_clip, btn_interrogate_booru = ui_sections.create_interrogate_buttons('control') + btn_interrogate = ui_sections.create_interrogate_button('control') with gr.Row(): input_buttons = [gr.Button('Select', visible=True, interactive=False), gr.Button('Inpaint', visible=True, interactive=True), gr.Button('Outpaint', visible=True, interactive=True)] with gr.Tab('Video', id='in-video') as tab_video: @@ -530,8 +530,7 @@ def create_ui(_blocks: gr.Blocks=None): input_type.change(fn=lambda x: gr.update(visible=x == 2), inputs=[input_type], outputs=[column_init]) btn_prompt_counter.click(fn=call_queue.wrap_queued_call(ui_common.update_token_counter), inputs=[prompt, steps], outputs=[prompt_counter]) btn_negative_counter.click(fn=call_queue.wrap_queued_call(ui_common.update_token_counter), inputs=[negative, steps], outputs=[negative_counter]) - btn_interrogate_clip.click(fn=helpers.interrogate_clip, inputs=[], outputs=[prompt]) - btn_interrogate_booru.click(fn=helpers.interrogate_booru, inputs=[], outputs=[prompt]) + btn_interrogate.click(fn=helpers.interrogate, inputs=[], outputs=[prompt]) select_fields = [input_mode, input_image, init_image, input_type, input_resize, input_inpaint, input_video, input_batch, input_folder] select_output = [output_tabs, preview_process, result_txt] diff --git a/modules/ui_control_helpers.py b/modules/ui_control_helpers.py index dae07b021..62688ea24 100644 --- a/modules/ui_control_helpers.py +++ b/modules/ui_control_helpers.py @@ -47,11 +47,21 @@ def initialize(): scripts.scripts_control.initialize_scripts(is_img2img=False, is_control=True) +def interrogate(): + prompt = None + try: + from modules.interrogate.interrogate import interrogate + prompt = interrogate(input_source[0]) + except Exception: + pass + return prompt + + def interrogate_clip(): # legacy function prompt = None try: - from modules.interrogate import legacy - prompt = legacy.interrogator.interrogate(input_source[0]) + from modules.interrogate import openclip + prompt = openclip.interrogator.interrogate(input_source[0]) except Exception: pass return gr.update() if prompt is None else prompt diff --git a/modules/ui_img2img.py b/modules/ui_img2img.py index 0b59e1b09..91113b03c 100644 --- a/modules/ui_img2img.py +++ b/modules/ui_img2img.py @@ -5,12 +5,13 @@ from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call from modules import timer, shared, ui_common, ui_sections, generation_parameters_copypaste, processing_vae -def process_interrogate(interrogation_function, mode, ii_input_files, ii_input_dir, ii_output_dir, *ii_singles): +def process_interrogate(mode, ii_input_files, ii_input_dir, ii_output_dir, *ii_singles): + from modules.interrogate.interrogate import interrogate mode = int(mode) if mode in {0, 1, 3, 4}: - return [interrogation_function(ii_singles[mode]), None] + return [interrogate(ii_singles[mode]), None] if mode == 2: - return [interrogation_function(ii_singles[mode]["image"]), None] + return [interrogate(ii_singles[mode]["image"]), None] if mode == 5: if len(ii_input_files) > 0: images = [f.name for f in ii_input_files] @@ -27,7 +28,7 @@ def process_interrogate(interrogation_function, mode, ii_input_files, ii_input_d img = Image.open(image) filename = os.path.basename(image) left, _ = os.path.splitext(filename) - print(interrogation_function(img), file=open(os.path.join(ii_output_dir, f"{left}.txt"), 'a', encoding='utf-8')) # pylint: disable=consider-using-with + print(interrogate(img), file=open(os.path.join(ii_output_dir, f"{left}.txt"), 'a', encoding='utf-8')) # pylint: disable=consider-using-with return [gr.update(), None] @@ -68,7 +69,7 @@ def create_ui(): state = gr.Textbox(value='', visible=False) with gr.TabItem('Image', id='img2img_image', elem_id="img2img_image_tab") as tab_img2img: img_init = gr.Image(label="", elem_id="img2img_image", show_label=False, source="upload", interactive=True, type="pil", tool="editor", image_mode="RGBA", height=512) - interrogate_clip, interrogate_booru = ui_sections.create_interrogate_buttons('img2img') + interrogate_btn = ui_sections.create_interrogate_button(tab='img2img') add_copy_image_controls('img2img', img_init) with gr.TabItem('Inpaint', id='img2img_inpaint', elem_id="img2img_inpaint_tab") as tab_inpaint: @@ -227,8 +228,7 @@ def create_ui(): ], outputs=[img2img_prompt, dummy_component], ) - interrogate_clip.click(fn=lambda *args: process_interrogate(ui_common.interrogate_clip, *args), **interrogate_args) - interrogate_booru.click(fn=lambda *args: process_interrogate(ui_common.interrogate_booru, *args), **interrogate_args) + interrogate_btn.click(fn=lambda *args: process_interrogate(*args), **interrogate_args) img2img_token_button.click(fn=wrap_queued_call(ui_common.update_token_counter), inputs=[img2img_prompt, steps], outputs=[img2img_token_counter]) img2img_negative_token_button.click(fn=wrap_queued_call(ui_common.update_token_counter), inputs=[img2img_negative_prompt, steps], outputs=[img2img_negative_token_counter]) diff --git a/modules/ui_postprocessing.py b/modules/ui_postprocessing.py index a3faf9766..903bc9e05 100644 --- a/modules/ui_postprocessing.py +++ b/modules/ui_postprocessing.py @@ -2,7 +2,7 @@ import json import gradio as gr 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 +from modules.interrogate import openclip def submit_info(image): @@ -46,31 +46,31 @@ 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') - 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') + clip_model = gr.Dropdown([], value=shared.opts.interrogate_clip_model, label='CLiP model') + ui_common.create_refresh_button(clip_model, openclip.refresh_clip_models, lambda: {"choices": openclip.refresh_clip_models()}, 'refresh_interrogate_models') + blip_model = gr.Dropdown(list(openclip.caption_models), value=shared.opts.interrogate_blip_model, label='Caption model') + mode = gr.Dropdown(openclip.caption_types, 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) + caption_max_length = gr.Number(label='Max length', value=shared.opts.interrogate_clip_max_length, 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=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=[]) + caption_max_length.change(fn=openclip.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[]) + chunk_size.change(fn=openclip.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[]) + min_flavors.change(fn=openclip.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[]) + max_flavors.change(fn=openclip.update_interrogate_params, inputs=[caption_max_length, chunk_size, min_flavors, max_flavors, flavor_intermediate_count], outputs=[]) + flavor_intermediate_count.change(fn=openclip.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(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) + btn_interrogate_img.click(openclip.interrogate_image, inputs=[image, clip_model, blip_model, mode], outputs=prompt) + btn_analyze_img.click(openclip.analyze_image, inputs=[image, clip_model, blip_model], outputs=[medium, artist, movement, trending, flavor]) + btn_unload.click(openclip.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) @@ -82,7 +82,7 @@ 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, legacy.get_clip_models, lambda: {"choices": legacy.get_clip_models()}, 'refresh_interrogate_models') + ui_common.create_refresh_button(clip_model, openclip.refresh_clip_models, lambda: {"choices": openclip.refresh_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"): @@ -90,11 +90,11 @@ def create_ui(): with gr.Row(): vqa_image = gr.Image(type='pil', label="Image") with gr.Row(): - vqa_question = gr.Textbox(label="Question", placeholder="Describe the image") + vqa_question = gr.Dropdown(label="Question", allow_custom_value=True, choices=vqa.vlm_prompts, value=vqa.vlm_prompts[2]) with gr.Row(): vqa_answer = gr.Textbox(label="Answer", lines=3) with gr.Row(elem_id='interrogate_buttons_query'): - vqa_model = gr.Dropdown(list(vqa.MODELS), value='MS Florence 2 Base', label='VQA Model') + vqa_model = gr.Dropdown(list(vqa.vlm_models), value=list(vqa.vlm_models)[0], label='VLM Model') vqa_submit = gr.Button("Interrogate", elem_id="interrogate_btn_interrogate", variant='primary') vqa_submit.click(vqa.interrogate, inputs=[vqa_question, vqa_image, vqa_model], outputs=[vqa_answer]) @@ -149,7 +149,7 @@ def create_ui(): ] ) btn_interrogate_batch.click( - fn=legacy.interrogate_batch, + fn=openclip.interrogate_batch, inputs=[batch_files, batch_folder, batch_str, clip_model, blip_model, mode, save_output], outputs=[batch], ) diff --git a/modules/ui_sections.py b/modules/ui_sections.py index fca9466b8..b55a30110 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -91,9 +91,11 @@ def create_resolution_inputs(tab): return width, height -def create_interrogate_button(tab: str, inputs: list, outputs: str): +def create_interrogate_button(tab: str, inputs: list = None, outputs: str = None): button_interrogate = gr.Button(ui_symbols.interrogate, elem_id=f"{tab}_interrogate", elem_classes=['interrogate']) - button_interrogate.click(fn=interrogate.interrogate, inputs=inputs, outputs=[outputs]) + if inputs is not None and outputs is not None: + button_interrogate.click(fn=interrogate.interrogate, inputs=inputs, outputs=[outputs]) + return button_interrogate def create_interrogate_buttons(tab): # legacy function diff --git a/repositories/blip/models/blip.py b/repositories/blip/models/blip.py index 32cdee3dc..83af0d11a 100644 --- a/repositories/blip/models/blip.py +++ b/repositories/blip/models/blip.py @@ -233,6 +233,4 @@ def load_checkpoint(model,url_or_filename): del state_dict[key] msg = model.load_state_dict(state_dict,strict=False) - print('load checkpoint from %s'%url_or_filename) return model,msg - diff --git a/requirements.txt b/requirements.txt index 14612bc63..e1995df7a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,7 +20,6 @@ piexif psutil pyyaml resize-right -rich toml voluptuous yapf @@ -32,6 +31,7 @@ invisible-watermark pi-heif # versioned +rich==13.9.4 safetensors==0.5.2 tensordict==0.1.2 peft==0.14.0 diff --git a/scripts/loopback.py b/scripts/loopback.py index a1541a161..99111d486 100644 --- a/scripts/loopback.py +++ b/scripts/loopback.py @@ -90,8 +90,8 @@ class Script(scripts.Script): if append_interrogation != "None": p.prompt = f"{original_prompt}, " if original_prompt else "" if append_interrogation == "CLIP": - from modules.interrogate import legacy - p.prompt += legacy.interrogator.interrogate(p.init_images[0]) + from modules.interrogate import openclip + p.prompt += openclip.interrogator.interrogate(p.init_images[0]) elif append_interrogation == "DeepBooru": from modules.interrogate import deepbooru p.prompt += deepbooru.model.tag(p.init_images[0])