From 8bf7ddffe34855c03d878573d770a6e1ce5d8925 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 14 May 2026 13:16:26 +0200 Subject: [PATCH] caption: add image analyze Signed-off-by: Vladimir Mandic --- .github/skills/check-api/SKILL.md | 12 +- .github/skills/update-docs/SKILL.md | 6 +- CHANGELOG.md | 6 + data/reference.json | 2 +- javascript/ui.js | 5 + modules/api/caption.py | 72 ++++++++++-- modules/api/validate.py | 1 + modules/caption/models_def.py | 39 +++++++ modules/caption/openclip.py | 21 ++-- modules/caption/vqa.py | 169 +++++++++++++++++++--------- modules/gr_hijack.py | 58 ++++++---- modules/sd_models.py | 6 +- modules/ui_caption.py | 156 ++++++++++++++++++++++--- 13 files changed, 432 insertions(+), 121 deletions(-) diff --git a/.github/skills/check-api/SKILL.md b/.github/skills/check-api/SKILL.md index 51b09c5bb..174e23ce8 100644 --- a/.github/skills/check-api/SKILL.md +++ b/.github/skills/check-api/SKILL.md @@ -35,17 +35,17 @@ This file is the route registration hub and must be treated as the source of tru For every endpoint, verify in this order: 1. Route validation: - - Route method and path are valid and unique after subpath handling. + - Route method and path are valid and unique after subpath handling. 2. Handler validation: - - Handler call signature is compatible with the route declaration. + - Handler call signature is compatible with the route declaration. 3. Request signature validation: - - Request body or query params implied by handler type hints are consistent with expected client usage. + - Request body or query params implied by handler type hints are consistent with expected client usage. 4. Response signature validation: - - Declared `response_model` is coherent with returned payload shape. + - Declared `response_model` is coherent with returned payload shape. 5. Auth validation: - - Authentication behavior is intentional (`auth=True` default in `add_api_route`). + - Authentication behavior is intentional (`auth=True` default in `add_api_route`). 6. OpenAPI validation: - - OpenAPI schema exposure is correct (including trailing-slash duplicate suppression). + - OpenAPI schema exposure is correct (including trailing-slash duplicate suppression). ## Procedure diff --git a/.github/skills/update-docs/SKILL.md b/.github/skills/update-docs/SKILL.md index dfac7ec5d..137dc4554 100644 --- a/.github/skills/update-docs/SKILL.md +++ b/.github/skills/update-docs/SKILL.md @@ -76,9 +76,9 @@ If targets are missing, ask for paths before editing. Extract from user prompt: - desired depth mode: - - syntax-only: fix markdown syntax/rendering issues only; do not rewrite wording or structure beyond what syntax requires - - readability: include syntax fixes plus clarity and scanability edits without broad restructuring - - full pass: include syntax, readability, structure normalization, terminology consistency, and broader doc cleanup + - syntax-only: fix markdown syntax/rendering issues only; do not rewrite wording or structure beyond what syntax requires + - readability: include syntax fixes plus clarity and scanability edits without broad restructuring + - full pass: include syntax, readability, structure normalization, terminology consistency, and broader doc cleanup If depth is missing, default to readability and state that assumption. diff --git a/CHANGELOG.md b/CHANGELOG.md index c2af8d502..81c6059f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,12 +2,18 @@ ## Update for 2026-05-14 +- **Features** + - **Captioning** new feature: analyze existing images for prompt adherence + *tip*: image analysis requires larger VLM model to produce quality output + new api endpoint: `/sdapi/v1/analyze` - **AI** - Cognitive analysis and improvements to *all* AI prompts - Automated fixes using `/check-` skills - Automated syntax, spelling and readability improvements to `/wiki` pages - **Fixes** - *hidream-o1* prequant loading + - `gradio` initial hijack + - `SmolVLM` captioning ## Update for 2026-05-13 diff --git a/data/reference.json b/data/reference.json index 347f3097e..d3cf7a4ad 100644 --- a/data/reference.json +++ b/data/reference.json @@ -951,4 +951,4 @@ "extras": "sampler: Default", "date": "2026 April" } -} \ No newline at end of file +} diff --git a/javascript/ui.js b/javascript/ui.js index 1cf6ee42f..b1e0c5078 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -252,6 +252,11 @@ function create_submit_args(args) { return res; } +function getCaptionActiveTab(...args) { + const res = create_tab_index_args('mode_caption', args); + return res; +} + function showSubmitButtons(tabname, show) {} function clearGallery(tabname) { diff --git a/modules/api/caption.py b/modules/api/caption.py index ab12c1920..4c133adfa 100644 --- a/modules/api/caption.py +++ b/modules/api/caption.py @@ -8,7 +8,7 @@ Provides three specialized backends and one unified dispatch endpoint: - POST /sdapi/v1/vqa — Vision-Language Models (Qwen, Gemma, Florence, Moondream, etc.) **Dispatch endpoint** (discriminated union routed by ``backend`` field): -- POST /sdapi/v1/caption — Routes to any backend via ``backend: "openclip" | "tagger" | "vlm"`` +- POST /sdapi/v1/caption — Routes to any backend via ``backend: "openclip" | "tagger" | "vlm" | "analyze"`` **Discovery endpoints** (GET, no request body): - GET /sdapi/v1/openclip — List available OpenCLIP models @@ -22,7 +22,7 @@ The dispatch endpoint uses a discriminated union (ReqCaptionDispatch) and a supe response model (ResCaptionDispatch) that includes fields from all backends. Core processing logic is shared between direct and dispatch handlers via -``do_openclip``, ``do_tagger``, and ``do_vqa`` functions to avoid duplication. +``do_openclip``, ``do_tagger``, and ``do_caption`` functions to avoid duplication. """ import threading @@ -212,9 +212,13 @@ class ReqCaptionVLM(BaseModel): keep_prefill: bool | None = Field(default=None, title="Keep Prefill", description="Keep prefill text in final output.") +class ReqCaptionAnalyze(ReqCaptionVLM): + backend: Literal["analyze"] = Field(..., description="Backend selector. Use 'analyze' for detailed image analysis using VLM.") + + # Discriminated union for the dispatch endpoint ReqCaptionDispatch = Annotated[ - ReqCaptionOpenCLIP | ReqCaptionTagger | ReqCaptionVLM, + ReqCaptionOpenCLIP | ReqCaptionTagger | ReqCaptionVLM | ReqCaptionAnalyze, Field(discriminator="backend") ] @@ -225,7 +229,7 @@ class ResCaptionDispatch(BaseModel): Contains fields from all backends - only relevant fields are populated based on the backend used. """ # Common - backend: str = Field(title="Backend", description="The backend that processed the request: 'openclip', 'tagger', or 'vlm'.") + backend: str = Field(title="Backend", description="The backend that processed the request: 'openclip', 'tagger', 'vlm', or 'analyze'.") # OpenCLIP fields caption: str | None = Field(default=None, title="Caption", description="Generated caption (OpenCLIP backend).") medium: str | None = Field(default=None, title="Medium", description="Detected artistic medium (OpenCLIP with analyze=True).") @@ -310,7 +314,7 @@ def build_vqa_kwargs(req) -> dict: return kwargs or None -def do_vqa(image, req): +def do_caption(image, req): """Core VLM captioning logic shared by direct and dispatch endpoints. Returns (answer, annotated_b64). @@ -336,6 +340,43 @@ def do_vqa(image, req): return answer, annotated_b64 +def do_analyze(image, req): + from modules.caption import vqa + from modules.caption.models_def import analyze_question + if req.question is None or len(req.question.strip()) < 2: + question = analyze_question + else: + question = req.question.strip() + if req.prompt is None or len(req.prompt.strip()) < 2: + from modules import images, infotext + info, _items = images.read_info_from_image(image) + items = infotext.parse(info) + prompt = (items.get('Prompt', None) or items.get('prompt', None)) if isinstance(items, dict) else None + if prompt is None: + return 'Error: No prompt found in image metadata.', None + else: + prompt = req.prompt.strip() + prompt = f"{question}\n\nDESCRIPTION: {prompt}" + answer = vqa.analyze( + question="Use Prompt", + system_prompt=req.system, + prompt=prompt, + image=image, + model_name=req.model, + prefill=req.prefill, + thinking_mode=req.thinking_mode, + generation_kwargs=build_vqa_kwargs(req) + ) + if isinstance(answer, str) and answer.startswith('Error:'): + raise HTTPException(status_code=422, detail=answer) + annotated_b64 = None + if req.include_annotated: + annotated_img = vqa.get_last_annotated_image() + if annotated_img is not None: + annotated_b64 = helpers.encode_pil_to_base64(annotated_img) + return answer, annotated_b64 + + def parse_tagger_scores(tags: str) -> dict: """Parse confidence scores from tagger output string.""" scores = {} @@ -496,7 +537,13 @@ def post_vqa(req: ReqVQA): - ``422``: Model returned an error (e.g., unsupported task for model) """ image = validate_image(req.image) - answer, annotated_b64 = do_vqa(image, req) + answer, annotated_b64 = do_caption(image, req) + return ResVQA(answer=answer, annotated_image=annotated_b64) + + +def post_analyze(req: ReqVQA): + image = validate_image(req.image) + answer, annotated_b64 = do_analyze(image, req) return ResVQA(answer=answer, annotated_image=annotated_b64) @@ -518,10 +565,14 @@ def post_caption_dispatch(req: ReqCaptionDispatch): WaifuDiffusion or DeepBooru anime/illustration tagging. Response populates ``tags`` (and ``scores`` when ``show_scores=True``). - 3. **VLM** (``backend: "vlm"``): + 3. **VLM** (``backend: "vlm"``): Vision-Language Models for flexible image understanding. Response populates ``answer`` (and ``annotated_image`` when ``include_annotated=True`` with detection tasks). + 4. **Analyze** (``backend: "analyze"``): + VLM-powered prompt analysis path that extracts or uses supplied prompt text and returns + an analysis response in ``answer`` (and ``annotated_image`` when available). + **Direct Endpoints** (backend-specific models, simpler interface): - POST /sdapi/v1/openclip — OpenCLIP only - POST /sdapi/v1/tagger — Tagger only @@ -542,8 +593,12 @@ def post_caption_dispatch(req: ReqCaptionDispatch): return ResCaptionDispatch(backend="tagger", tags=tags, scores=scores) elif req.backend == "vlm": image = validate_image(req.image) - answer, annotated_b64 = do_vqa(image, req) + answer, annotated_b64 = do_caption(image, req) return ResCaptionDispatch(backend="vlm", answer=answer, annotated_image=annotated_b64) + elif req.backend == 'analyze': + image = validate_image(req.image) + answer, annotated_b64 = do_analyze(image, req) + return ResCaptionDispatch(backend="analyze", answer=answer, annotated_image=annotated_b64) else: raise HTTPException(status_code=400, detail=f"Unknown backend: {req.backend}") @@ -667,6 +722,7 @@ def register_api(api): api.add_api_route("/sdapi/v1/caption", post_caption_dispatch, methods=["POST"], response_model=ResCaptionDispatch, tags=["Caption"]) api.add_api_route("/sdapi/v1/openclip", post_caption, methods=["POST"], response_model=ResCaption, tags=["Caption"]) api.add_api_route("/sdapi/v1/vqa", post_vqa, methods=["POST"], response_model=ResVQA, tags=["Caption"]) + api.add_api_route("/sdapi/v1/analyze", post_analyze, methods=["POST"], response_model=ResVQA, tags=["Caption"]) api.add_api_route("/sdapi/v1/vqa/models", get_vqa_models, methods=["GET"], response_model=list[ItemVLMModel], tags=["Caption"]) api.add_api_route("/sdapi/v1/vqa/prompts", get_vqa_prompts, methods=["GET"], response_model=ResVLMPrompts, tags=["Caption"]) api.add_api_route("/sdapi/v1/tagger", post_tagger, methods=["POST"], response_model=ResTagger, tags=["Caption"]) diff --git a/modules/api/validate.py b/modules/api/validate.py index 1ff77be75..bd2eb920f 100644 --- a/modules/api/validate.py +++ b/modules/api/validate.py @@ -31,6 +31,7 @@ log_cost = { "/sdapi/v1/memory": -1, "/sdapi/v1/platform": -1, "/sdapi/v1/checkpoint": -1, + "/sdapi/v1/gpu-smi": -1, "/sdapi/v1/status": 60, "/sdapi/v1/progress": 60, } diff --git a/modules/caption/models_def.py b/modules/caption/models_def.py index 475fb5a0d..21ecbc410 100644 --- a/modules/caption/models_def.py +++ b/modules/caption/models_def.py @@ -89,6 +89,7 @@ vlm_default = "Alibaba Qwen 2.5 VL 3B" # Default system prompt vlm_system = 'You are image captioning expert, creative, unbiased and uncensored.' +vlm_analyze = 'You are an image analysis expert with focus on details, unbiased and uncensored.' # Common prompts (work with all VLM models) vlm_prompts_common = [ @@ -174,6 +175,44 @@ vlm_prompts = vlm_prompts_common + vlm_prompts_florence + vlm_prompts_promptgen vlm_prefill = 'Answer: the image shows' +analyze_question = """Compare the image against the provided DESCRIPTION. + +**Instructions**: +- Compare IMAGE vs DESCRIPTION using only visually verifiable content. +- Use only details that can be directly seen in the image. +- Ignore non-visual or unverifiable claims in DESCRIPTION, including platform/source references (for example Instagram), camera/device/lens/settings, resolution/quality tags (for example 4k, UHD), style buzzwords (for example inspirational), and subjective attractiveness claims (for example beautiful, stunning). +- Never mark as Missing or Differences any subjective person descriptors or archetypes (for example young, beautiful, goddess, handsome, elegant, sexy, heroic) unless they are replaced by objective visual traits. +- Treat implicit, inferred, interpretive, or hedged wording as non-actionable (for example implied mood/time, seems, appears, suggests) and do not mark it as Missing or Differences. +- If image has visible flaws or artifacts, note them in Flaws. +- If subjects or objects in the image have visibly incorrect or inconsistent details (for example extra limbs, distorted faces, impossible anatomy), note them in Flaws. +- Do not report what DESCRIPTION says; compare image content only. + +**Output format** (plain text only; use these sections in this order when they have content): +Matching: +- +Missing: +- +Extras: +- +Differences: +- +Flaws: +- +Summary: +- <1-2 sentences on overall alignment and key gap> +- Score: <0.0-1.0 alignment score> +- <1-2 sentences on overall image quality based on clarity/composition and visual interest> +- Quality: <0.0-1.0 visual quality score based only on image clarity/composition> + +**Rules**: +- All sections must be a bullet list with max 5 bullets per section. +- Keep each bullet to one short sentence. +- Omit sections that have no items. +- Only include Missing items that are explicit, concrete, and directly verifiable from image content. +- Do not output chain-of-thought, reasoning steps, or meta commentary. +- Do not use phrases like "the description mentions" or "the prompt says" or prefixes like "A", "The", etc. in bullets; start directly with the detail. +""" + def get_vlm_repo(display_name: str) -> str: """Look up repo ID from display name, stripping any trailing symbols.""" diff --git a/modules/caption/openclip.py b/modules/caption/openclip.py index 0840514dc..a7e3134a3 100644 --- a/modules/caption/openclip.py +++ b/modules/caption/openclip.py @@ -11,9 +11,7 @@ from modules.logger import log, console debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None - -# Per-request overrides for API calls -_clip_overrides = None +_clip_overrides = None # Per-request overrides for API calls def get_clip_setting(name): @@ -97,7 +95,6 @@ def update_caption_params(): ci.caption_offload = shared.opts.caption_offload - def get_clip_models(): return clip_models @@ -186,12 +183,12 @@ def load_captioner(clip_model, blip_model): _apply_blip2_fix(ci.caption_model, ci.caption_processor) log.debug(f'CLIP load: time={time.time()-t0:.2f}') else: - debug_log(f'CLIP: models already loaded clip="{clip_model}" blip="{blip_model}"') + debug_log(f'CLIP load: clip="{clip_model}" blip="{blip_model}" already loaded') def unload_clip_model(): if ci is not None and shared.opts.caption_offload: - log.debug('CLIP unload: offloading models to CPU') + debug_log('CLIP unload: offloading models to CPU') # Direct .to() instead of sd_models.move_model — models are from clip_interrogator, not transformers if ci.caption_model is not None and hasattr(ci.caption_model, 'to'): ci.caption_model.to(devices.cpu) @@ -233,12 +230,14 @@ def caption(image, mode, base_caption=None): return prompt - def caption_image(image, clip_model, blip_model, mode, overrides=None): global _clip_overrides # pylint: disable=global-statement + if image is None: + log.error('CLIP: image=None') + return 'CLIP error: no image provided' jobid = shared.state.begin('Caption CLiP') t0 = time.time() - log.info(f'CLIP: mode="{mode}" clip="{clip_model}" blip="{blip_model}" image_size={image.size if image else None}') + log.info(f'CLIP caption: mode="{mode}" clip="{clip_model}" blip="{blip_model}" image_size={image.size if image else None}') if overrides: debug_log(f'CLIP: overrides={overrides}') try: @@ -256,9 +255,9 @@ def caption_image(image, clip_model, blip_model, mode, overrides=None): if shared.opts.caption_offload: unload_clip_model() devices.torch_gc() - log.debug(f'CLIP: complete time={time.time()-t0:.2f}') + log.debug(f'CLIP complete: time={time.time()-t0:.2f}') except Exception as e: - prompt = f"Exception {type(e)}" + prompt = f"CLIP error: {type(e)}" log.error(f'CLIP: {e}') errors.display(e, 'Caption') finally: @@ -268,7 +267,6 @@ def caption_image(image, clip_model, blip_model, mode, overrides=None): return prompt - def caption_batch(batch_files, batch_folder, batch_str, clip_model, blip_model, mode, write, append, recursive): files = [] if batch_files is not None: @@ -318,7 +316,6 @@ def caption_batch(batch_files, batch_folder, batch_str, clip_model, blip_model, return '\n\n'.join(prompts) - def analyze_image(image, clip_model, blip_model): t0 = time.time() log.info(f'CLIP analyze: clip="{clip_model}" blip="{blip_model}" image_size={image.size if image else None}') diff --git a/modules/caption/vqa.py b/modules/caption/vqa.py index 74cc3cdca..0e80ccbd2 100644 --- a/modules/caption/vqa.py +++ b/modules/caption/vqa.py @@ -13,16 +13,49 @@ from modules import shared, devices, errors, model_quant, sd_models, sd_models_c from modules.sd_offload_aux import register_aux, deregister_aux, move_aux_to_gpu, offload_aux from modules.logger import log, console from modules.caption import vqa_detection -from modules.caption.models_def import vlm_models, vlm_system, vlm_default, vlm_prefill, vlm_prompts, vlm_prompt_mapping, vlm_prompt_placeholders, vlm_prompts_common, vlm_prompts_florence, vlm_prompts_moondream, vlm_prompts_moondream2, vlm_prompts_promptgen, get_vlm_repo +from modules.caption.models_def import vlm_models, vlm_system, vlm_analyze, vlm_default, vlm_prefill, vlm_prompts, vlm_prompt_mapping, vlm_prompt_placeholders, vlm_prompts_common, vlm_prompts_florence, vlm_prompts_moondream, vlm_prompts_moondream2, vlm_prompts_promptgen, analyze_question, get_vlm_repo # pylint: disable=unused-import # Debug logging - function-based to avoid circular import debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None + def debug(*args, **kwargs): if debug_enabled: log.trace(*args, **kwargs) +class BatchWriter: + def __init__(self, folder, mode='w', fmt='txt', filename=None): + self.folder = folder + self.file = None + self.mode = mode + self.format = fmt + self.entries = [] + self.filename = filename + + def add(self, file, text): + if self.format == 'txt': + txt_file = os.path.splitext(file)[0] + ".txt" if not self.filename else self.filename + if self.mode == 'a': + text = '\n' + text + with open(os.path.join(self.folder, txt_file), self.mode, encoding='utf-8') as f: + f.write(text) + if self.format == 'json': + json_file = os.path.splitext(file)[0] + ".json" if not self.filename else self.filename + self.mode = 'w' + try: + dct = json.loads(text) + entry = {"file": file, **dct} + except Exception: + entry = {"file": file, "answer": text} + self.entries.append(entry) + with open(os.path.join(self.folder, json_file), self.mode, encoding='utf-8') as f: + f.write(json.dumps(self.entries, ensure_ascii=False, indent=2)) + + def close(self): + if self.file is not None: + self.file.close() + def get_prompts_for_model(model_name: str) -> list: """Get available prompts based on selected model.""" @@ -213,7 +246,8 @@ def clean(response, question, prefill=None): else: # Remove prefill if it's present in the cleaned response if len(prefill_text) > 0 and response.startswith(prefill_text): - response = response[len(prefill_text):].strip() + response = response[len(prefill_text):] + response = response.replace('\n\n', '\n').strip() return response @@ -759,7 +793,7 @@ class VQA: self.loaded = repo devices.torch_gc() - def _mistral(self, question: str, image: Image.Image, repo: str, system_prompt: str | None = None, model_name: str | None = None, prefill: str | None = None, thinking_mode: bool = False): + def _mistral(self, question: str, image: Image.Image, repo: str, system_prompt: str | None = None, model_name: str | None = None, prefill: str | None = None, thinking_mode: bool = False): # pylint: disable=unused-argument self._load_mistral(repo) move_aux_to_gpu('vqa') cls_name = self.model.__class__.__name__ @@ -911,7 +945,7 @@ class VQA: log.debug(f'Caption load: vlm="{repo}"') self._unload_current() quant_args = model_quant.create_config(module='LLM') - self.model = transformers.AutoModelForVision2Seq.from_pretrained( + self.model = transformers.AutoModelForImageTextToText.from_pretrained( repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=devices.dtype, @@ -1358,7 +1392,7 @@ class VQA: self.last_annotated_image = None self.last_detection_data = None self._generation_overrides = generation_kwargs # Set per-request overrides - jobid = shared.state.begin('Caption LLM') + jobid = shared.state.begin('VLM') t0 = time.time() model_name = model_name or shared.opts.caption_vlm_model prefill = vlm_prefill if prefill is None else prefill # Use provided prefill when specified @@ -1379,7 +1413,7 @@ class VQA: return 'Error: No input image provided. Please upload or select an image.' # Convert friendly prompt names to internal tokens/commands - if question == "Use Prompt": + if question.lower() == "use prompt": # Use content from Prompt field directly - requires user input if not prompt or len(prompt.strip()) < 2: log.error(f'VQA caption: model="{model_name}" error="Please enter a prompt"') @@ -1401,21 +1435,21 @@ class VQA: question = get_internal_prompt(question, prompt) # else: question is already an internal token or custom text - from modules import modelloader - modelloader.hf_login() - sd_models.set_caption_load_options() + if model_name is None: + log.error(f'Caption: type=vlm model="{model_name}" no model selected') + shared.state.end(jobid) + return '' + vqa_model = get_vlm_repo(model_name) + if vqa_model == model_name and model_name not in vlm_models.values(): + log.error(f'Caption: type=vlm model="{model_name}" unknown') + shared.state.end(jobid) + return '' + if self.model is None or self.loaded != vqa_model: + from modules import modelloader + modelloader.hf_login() + sd_models.set_caption_load_options() try: - if model_name is None: - log.error(f'Caption: type=vlm model="{model_name}" no model selected') - shared.state.end(jobid) - return '' - vqa_model = get_vlm_repo(model_name) - if vqa_model == model_name and model_name not in vlm_models.values(): - log.error(f'Caption: type=vlm model="{model_name}" unknown') - shared.state.end(jobid) - return '' - handler = 'unknown' if 'git' in vqa_model.lower(): handler = 'git' @@ -1519,26 +1553,43 @@ class VQA: shared.state.end(jobid) return answer + def analyze( + self, + question: str = "", + system_prompt: str | None = None, + prompt: str | None = None, + image: list[Image.Image] | Image.Image | dict | None = None, + model_name: str | None = None, + prefill: str | None = None, # pylint: disable=unused-argument + thinking_mode: bool | None = None, + quiet: bool = False, + generation_kwargs: dict | None = None, + ) -> str: + if question is None or len(question.strip()) < 2: + question = analyze_question + if prompt is None or len(prompt.strip()) < 2: + from modules import images, infotext + info, _items = images.read_info_from_image(image) + items = infotext.parse(info) + prompt = (items.get('Prompt', None) or items.get('prompt', None)) if isinstance(items, dict) else None + if prompt is None: + log.error('VQA analyze: no prompt found in image metadata') + return 'Error: No prompt found in image metadata.' + prompt = f"{question}\n\nDESCRIPTION: {prompt}" + answer = self.caption( + question="Use Prompt", + system_prompt=system_prompt, + prompt=prompt, + image=image, + model_name=model_name, + prefill='', + thinking_mode=thinking_mode, + quiet=quiet, + generation_kwargs=generation_kwargs, + ) + return answer - def batch(self, model_name, system_prompt, batch_files, batch_folder, batch_str, question, prompt, write, append, recursive, prefill=None, thinking_mode=False): - class BatchWriter: - def __init__(self, folder, mode='w'): - self.folder = folder - self.csv = None - self.file = None - self.mode = mode - - def add(self, file, prompt_text): - txt_file = os.path.splitext(file)[0] + ".txt" - if self.mode == 'a': - prompt_text = '\n' + prompt_text - with open(os.path.join(self.folder, txt_file), self.mode, encoding='utf-8') as f: - f.write(prompt_text) - - def close(self): - if self.file is not None: - self.file.close() - + def batch(self, model_name, system_prompt, batch_files, batch_folder, batch_str, question, prompt, save_txt, append_txt, save_json, recursive, prefill=None, thinking_mode=False, vlm_mode='caption'): files = [] if batch_files is not None: files += [f.name for f in batch_files] @@ -1552,34 +1603,47 @@ class VQA: return '' jobid = shared.state.begin('Caption batch') prompts = [] - if write: - mode = 'w' if not append else 'a' - writer = BatchWriter(os.path.dirname(files[0]), mode=mode) + if save_txt: + mode = 'w' if not append_txt else 'a' + writer_txt = BatchWriter(folder=os.path.dirname(files[0]), mode=mode, fmt='txt') + if save_json: + writer_json = BatchWriter(folder=os.path.dirname(files[0]), fmt='json', filename=f'{vlm_mode}.json') orig_offload = shared.opts.caption_offload shared.opts.caption_offload = False try: import rich.progress as rp - pbar = rp.Progress(rp.TextColumn('[cyan]Caption:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console) + pbar = rp.Progress(rp.TextColumn('[cyan]VLM:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console) with pbar: task = pbar.add_task(total=len(files), description='starting...') for file in files: - pbar.update(task, advance=1, description=file) + pbar.update(task, advance=1, description=f'file={file}') try: if shared.state.interrupted: break - img = Image.open(file) - result = self.caption(question, system_prompt, prompt, img, model_name, prefill, thinking_mode, quiet=True) - # Save annotated image if available - if self.last_annotated_image and write: + try: + img = Image.open(file) + except Exception: + continue + if vlm_mode == 'caption': + result = self.caption(question, system_prompt, prompt, img, model_name, prefill, thinking_mode, quiet=True) + elif vlm_mode == 'analyze': + result = self.analyze(question, system_prompt, prompt, img, model_name, prefill, thinking_mode, quiet=True) + else: + result = f'Unknown mode: {vlm_mode}' + if self.last_annotated_image and (save_txt or save_json): # save annotated image if available annotated_path = os.path.splitext(file)[0] + "_annotated.png" self.last_annotated_image.save(annotated_path) prompts.append(result) - if write: - writer.add(file, result) + if save_txt: + writer_txt.add(file, result) + if save_json: + writer_json.add(file, result) except Exception as e: log.error(f'Caption batch: {e}') - if write: - writer.close() + if save_txt: + writer_txt.close() + if save_json: + writer_json.close() finally: shared.opts.caption_offload = orig_offload offload_aux('vqa') @@ -1604,6 +1668,9 @@ def caption(*args, **kwargs): return get_instance().caption(*args, **kwargs) +def analyze(*args, **kwargs): + return get_instance().analyze(*args, **kwargs) + def unload_model(): return get_instance().unload() diff --git a/modules/gr_hijack.py b/modules/gr_hijack.py index 082ace269..5940bab32 100644 --- a/modules/gr_hijack.py +++ b/modules/gr_hijack.py @@ -1,6 +1,6 @@ import time from PIL import Image -import gradio as gr +import gradio import gradio.processing_utils from modules import scripts_manager, patches, gr_tempdir from modules.logger import log @@ -13,7 +13,7 @@ original_BlockContext_init = None original_Blocks_get_config_file = None -def process_kanvas(self, x): # only used when kanvas overrides gr.Image object +def process_kanvas(self, x): # only used when kanvas overrides gradio.Image object import numpy as np t0 = time.time() image_data = list(x.get('image', {}).values()) @@ -72,7 +72,7 @@ def gr_image_preprocess(self, x): def add_classes_to_gradio_component(comp): """ - this adds gradio-* to the component for css styling (ie gradio-button to gr.Button), as well as some others + this adds gradio-* to the component for css styling (ie gradio-button to gradio.Button), as well as some others """ comp.elem_classes = [f"gradio-{comp.get_block_name()}", *(comp.elem_classes or [])] if getattr(comp, 'multiselect', False): @@ -150,10 +150,26 @@ def reset_gradio_sessions(job_id): def patch_gradio(): - orig_cancel_tasks = gradio.utils.cancel_tasks - orig_restore_session_state = gradio.route_utils.restore_session_state - orig_call_prediction = gradio.queueing.Queue.call_prediction - orig_blocks_preprocess_data = gradio.blocks.Blocks.preprocess_data + try: + orig_cancel_tasks = gradio.utils.cancel_tasks + except Exception: + log.error(f'Gradio patch: version={gradio.__version__} cancel_tasks not found') + orig_cancel_tasks = None + try: + orig_restore_session_state = gradio.route_utils.restore_session_state + except Exception: + log.error(f'Gradio patch: version={gradio.__version__} restore_session_state not found') + orig_restore_session_state = None + try: + orig_call_prediction = gradio.queueing.Queue.call_prediction + except Exception: + log.error(f'Gradio patch: version={gradio.__version__} call_prediction not found') + orig_call_prediction = None + try: + orig_blocks_preprocess_data = gradio.blocks.Blocks.preprocess_data + except Exception: + log.error(f'Gradio patch: version={gradio.__version__} preprocess_data not found') + orig_blocks_preprocess_data = None async def wrap_cancel_tasks(task_ids: set[str]): log.error(f'Gradio cancel: task={task_ids}') @@ -208,10 +224,14 @@ def patch_gradio(): log.error(f"Gradio preprocess: {e}") raise - gradio.queueing.Queue.call_prediction = wrap_call_prediction - gradio.route_utils.restore_session_state = wrap_restore_session_state - gradio.utils.cancel_tasks = wrap_cancel_tasks - gradio.blocks.Blocks.preprocess_data = wrap_blocks_preprocess_data + if orig_call_prediction is not None: + gradio.queueing.Queue.call_prediction = wrap_call_prediction + if orig_restore_session_state is not None: + gradio.route_utils.restore_session_state = wrap_restore_session_state + if orig_cancel_tasks is not None: + gradio.utils.cancel_tasks = wrap_cancel_tasks + if orig_blocks_preprocess_data is not None: + gradio.blocks.Blocks.preprocess_data = wrap_blocks_preprocess_data def patch_gradio_future(): @@ -248,14 +268,14 @@ def init(): global hijacked, original_IOComponent_init, original_Block_get_config, original_BlockContext_init, original_Blocks_get_config_file # pylint: disable=global-statement if hijacked: return - gr.components.Image.preprocess = gr_image_preprocess - if hasattr(gr.components, 'IOComponent'): - gr.components.IOComponent.pil_to_temp_file = gr_tempdir.pil_to_temp_file - original_IOComponent_init = patches.patch(__name__, obj=gr.components.IOComponent, field="__init__", replacement=IOComponent_init) - original_Block_get_config = patches.patch(__name__, obj=gr.blocks.Block, field="get_config", replacement=Block_get_config) - original_BlockContext_init = patches.patch(__name__, obj=gr.blocks.BlockContext, field="__init__", replacement=BlockContext_init) - original_Blocks_get_config_file = patches.patch(__name__, obj=gr.blocks.Blocks, field="get_config_file", replacement=Blocks_get_config_file) + gradio.components.Image.preprocess = gr_image_preprocess + if hasattr(gradio.components, 'IOComponent'): + gradio.components.IOComponent.pil_to_temp_file = gr_tempdir.pil_to_temp_file + original_IOComponent_init = patches.patch(__name__, obj=gradio.components.IOComponent, field="__init__", replacement=IOComponent_init) + original_Block_get_config = patches.patch(__name__, obj=gradio.blocks.Block, field="get_config", replacement=Block_get_config) + original_BlockContext_init = patches.patch(__name__, obj=gradio.blocks.BlockContext, field="__init__", replacement=BlockContext_init) + original_Blocks_get_config_file = patches.patch(__name__, obj=gradio.blocks.Blocks, field="get_config_file", replacement=Blocks_get_config_file) patch_gradio() - if not gr.__version__.startswith('3.43'): + if not gradio.__version__.startswith('3.43'): patch_gradio_future() hijacked = True diff --git a/modules/sd_models.py b/modules/sd_models.py index 288d10ab3..004a75c91 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -84,11 +84,11 @@ def set_caption_load_options(): else: sd_hijack_accelerate.restore_accelerate() if (shared.opts.runai_streamer_diffusers or shared.opts.runai_streamer_transformers) and (sys.platform == 'linux'): - log.debug(f'Caption loader: to_gpu={shared.opts.caption_to_gpu} runai chunk={os.environ.get("RUNAI_STREAMER_CHUNK_BYTESIZE", "N/A")} limit={os.environ.get("RUNAI_STREAMER_MEMORY_LIMIT", "N/A")}') + log.debug(f'Caption loader: gpu={shared.opts.caption_to_gpu} runai=True chunk={os.environ.get("RUNAI_STREAMER_CHUNK_BYTESIZE", "N/A")} limit={os.environ.get("RUNAI_STREAMER_MEMORY_LIMIT", "N/A")}') sd_hijack_safetensors.hijack_safetensors(shared.opts.runai_streamer_diffusers, shared.opts.runai_streamer_transformers) else: if shared.opts.caption_to_gpu: - log.debug(f'Caption loader: to_gpu={shared.opts.caption_to_gpu}') + log.debug(f'Caption loader: gpu={shared.opts.caption_to_gpu}') sd_hijack_safetensors.restore_safetensors() sd_hijack_hfhub.init_hijack() @@ -744,7 +744,7 @@ def load_sdnq_model(checkpoint_info, pipeline, diffusers_load_config, op): module, name, t = load_sdnq_module(checkpoint_info.path, module_name, load_method=load_method) if module is not None: modules[name] = module - log.debug(f'Load {op}: module="{checkpoint_info.name}" module="{name}" direct={shared.opts.diffusers_to_gpu} prequant=sdnq method={load_method} time={t:.2f}') + log.debug(f'Load {op}: module="{checkpoint_info.name}" module="{name}" gpu={shared.opts.diffusers_to_gpu} prequant=sdnq method={load_method} time={t:.2f}') """ futures = [] diff --git a/modules/ui_caption.py b/modules/ui_caption.py index 5d4693e61..3501496eb 100644 --- a/modules/ui_caption.py +++ b/modules/ui_caption.py @@ -6,14 +6,38 @@ from modules.caption import openclip default_task = "Normal Caption" -def vlm_caption_wrapper(question, system_prompt, prompt, image, model_name, prefill, thinking_mode): - """Wrapper for vqa.caption that handles annotated image display.""" - from modules.caption import vqa - answer = vqa.caption(question, system_prompt, prompt, image, model_name, prefill, thinking_mode) - annotated_image = vqa.get_last_annotated_image() - if annotated_image is not None: - return answer, gr.update(value=annotated_image, visible=True) - return answer, gr.update(visible=False) + +def caption_wrapper(tab, image, + vlm_question, vlm_system, vlm_prompt, vlm_model, vlm_prefill, vlm_thinking_mode, + analyze_question, analyze_system, analyze_prompt, analyze_model, + clip_model, blip_model, clip_mode, + wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape, + ): + """Wrapper for vqa.caption, vqa.analysis, openclip.caption_image, tagger.tag.""" + if tab <= 0: + log.debug('Caption: mode="VLM Caption"') + from modules.caption import vqa + answer = vqa.caption(vlm_question, vlm_system, vlm_prompt, image, vlm_model, vlm_prefill, vlm_thinking_mode) + annotated_image = vqa.get_last_annotated_image() + if annotated_image is not None: + return answer, gr.update(value=annotated_image, visible=True) + return answer, gr.update(visible=False) + elif tab == 1: + log.debug('Caption: mode="VLM Analyze"') + from modules.caption import vqa + answer = vqa.analyze(analyze_question, analyze_system, analyze_prompt, image, analyze_model, vlm_thinking_mode) + return answer, gr.update(visible=False) + elif tab == 2: + log.debug('Caption: mode="OpenCLIP"') + caption = openclip.caption_image(image, clip_model, blip_model, clip_mode) + return caption, gr.update(visible=False) + elif tab == 3: + log.debug('Caption: mode="Tagger"') + tags = tagger_tag_wrapper(image, wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape) + return tags, gr.update(visible=False) + else: + log.error(f'Caption: mode={tab} unknown') + return 'Unknown caption mode', gr.update(visible=False) def update_vlm_prompts_for_model(model_name): @@ -153,6 +177,7 @@ def create_ui(): with gr.Row(): image = gr.Image(type='pil', label="Image", height=512, visible=True, image_mode='RGB', elem_id='caption_image') with gr.Tabs(elem_id="mode_caption"): + with gr.Tab("VLM Caption", elem_id="tab_vlm_caption"): from modules.caption import vqa current_vlm_model = shared.opts.caption_vlm_model or vqa.vlm_default @@ -201,13 +226,47 @@ def create_ui(): with gr.Row(): vlm_batch_str = gr.Textbox(label="Folder", value="", interactive=True, elem_id='vlm_batch_str') with gr.Row(): - vlm_save_output = gr.Checkbox(label='Save Caption Files', value=True, elem_id="vlm_save_output") + vlm_save_txt = gr.Checkbox(label='Save Caption Files', value=True, elem_id="vlm_save_output") vlm_save_append = gr.Checkbox(label='Append Caption Files', value=False, elem_id="vlm_save_append") + vlm_save_json = gr.Checkbox(label='Save Caption JSON', value=True, elem_id="vlm_save_json") vlm_folder_recursive = gr.Checkbox(label='Recursive', value=False, elem_id="vlm_folder_recursive") with gr.Row(elem_id='caption_buttons_batch'): btn_vlm_caption_batch = gr.Button("Batch Caption", variant='primary', elem_id="btn_vlm_caption_batch") with gr.Row(): btn_vlm_caption = gr.Button("Caption", variant='primary', elem_id="btn_vlm_caption") + + with gr.Tab("VLM Analyze", elem_id="tab_vlm_analyze"): + from modules.caption import vqa + analyze_question_placeholder = 'Enter your analysis question or leave blank to use default' + analyze_prompt_placeholder = 'Enter your prompt to match with image or leave blank to use image metadata' + with gr.Row(): + analyze_system = gr.Textbox(label="System Prompt", value=vqa.vlm_analyze, lines=1, elem_id='analyze_system') + with gr.Row(): + analyze_question = gr.Textbox(label="Question", placeholder=analyze_question_placeholder, lines=1, elem_id='analyze_question') + with gr.Row(): + analyze_prompt = gr.Textbox(label="Prompt", placeholder=analyze_prompt_placeholder, lines=2, elem_id='analyze_prompt') + with gr.Row(elem_id='caption_buttons_query'): + analyze_model = gr.Dropdown(list(vqa.vlm_models), value=current_vlm_model, label='VLM Model', elem_id='analyze_model') + with gr.Row(): + analyze_load_btn = gr.Button(value='Load', elem_id='analyze_load', variant='secondary') + analyze_unload_btn = gr.Button(value='Unload', elem_id='analyze_unload', variant='secondary') + with gr.Accordion(label='Analyze: Batch', open=False, visible=True): + with gr.Row(): + analyze_batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], interactive=True, height=100, elem_id='analyze_batch_files') + with gr.Row(): + analyze_batch_folder = gr.File(label="Folder", show_label=True, file_count='directory', file_types=['image'], interactive=True, height=100, elem_id='analyze_batch_folder') + with gr.Row(): + analyze_batch_str = gr.Textbox(label="Folder", value="", interactive=True, elem_id='analyze_batch_str') + with gr.Row(): + analyze_save_txt = gr.Checkbox(label='Save Analysis Files', value=True, elem_id="analyze_save_txt") + analyze_save_append = gr.Checkbox(label='Append Analysis Files', value=False, elem_id="analyze_save_append") + analyze_save_json = gr.Checkbox(label='Save Analysis JSON', value=False, elem_id="analyze_save_json") + analyze_folder_recursive = gr.Checkbox(label='Recursive', value=False, elem_id="analyze_folder_recursive") + with gr.Row(elem_id='caption_buttons_batch'): + btn_analyze_caption_batch = gr.Button("Batch Analysis", variant='primary', elem_id="btn_analyze_caption_batch") + with gr.Row(): + btn_analyze_caption = gr.Button("Analyze", variant='primary', elem_id="btn_analyze_caption") + with gr.Tab("OpenCLiP", elem_id='tab_openclip'): with gr.Row(): clip_model = gr.Dropdown([], value=shared.opts.caption_openclip_model, label='CLiP Model', elem_id='clip_clip_model') @@ -244,8 +303,9 @@ def create_ui(): with gr.Row(): btn_clip_caption_batch = gr.Button("Batch Caption", variant='primary', elem_id="btn_clip_caption_batch") with gr.Row(): - btn_clip_caption_img = gr.Button("Caption", variant='primary', elem_id="btn_clip_caption_img") - btn_clip_analyze_img = gr.Button("Analyze", variant='primary', elem_id="btn_clip_analyze_img") + btn_clip_caption_img = gr.Button("CLiP Caption", variant='primary', elem_id="btn_clip_caption_img") + btn_clip_analyze_img = gr.Button("CLiP Analyze", variant='primary', elem_id="btn_clip_analyze_img") + with gr.Tab("Tagger", elem_id='tab_tagger'): from modules.caption import tagger with gr.Row(): @@ -294,6 +354,7 @@ def create_ui(): label="Default Caption Type", elem_id="default_caption_type" ) + with gr.Column(variant='compact', elem_id='caption_output'): with gr.Row(elem_id='caption_output_prompt'): prompt = gr.Textbox(label="Answer", lines=12, placeholder="ai generated image description") @@ -309,13 +370,69 @@ def create_ui(): with gr.Row(elem_id='copy_buttons_caption'): copy_caption_buttons = generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "control", "extras"]) - btn_clip_caption_img.click(openclip.caption_image, inputs=[image, clip_model, blip_model, clip_mode], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) - btn_clip_analyze_img.click(openclip.analyze_image, inputs=[image, clip_model, blip_model], outputs=[medium, artist, movement, trending, flavor, clip_labels_text]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) - btn_clip_caption_batch.click(fn=openclip.caption_batch, inputs=[clip_batch_files, clip_batch_folder, clip_batch_str, clip_model, blip_model, clip_mode, clip_save_output, clip_save_append, clip_folder_recursive], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) - btn_vlm_caption.click(fn=vlm_caption_wrapper, inputs=[vlm_question, vlm_system, vlm_prompt, image, vlm_model, vlm_prefill, vlm_thinking_mode], outputs=[prompt, output_image]) - btn_vlm_caption_batch.click(fn=vqa.batch, inputs=[vlm_model, vlm_system, vlm_batch_files, vlm_batch_folder, vlm_batch_str, vlm_question, vlm_prompt, vlm_save_output, vlm_save_append, vlm_folder_recursive, vlm_prefill, vlm_thinking_mode], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) - btn_wd_tag.click(fn=tagger_tag_wrapper, inputs=[image, wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) - btn_wd_tag_batch.click(fn=tagger_batch_wrapper, inputs=[wd_model, wd_batch_files, wd_batch_folder, wd_batch_str, wd_save_output, wd_save_append, wd_folder_recursive, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape], outputs=[prompt]).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) + dummy = gr.Label(value='-1', visible=False) + btn_vlm_caption.click( + _js="getCaptionActiveTab", # js to insert current tab name as first argument + fn=caption_wrapper, + inputs=[dummy, image, + vlm_question, vlm_system, vlm_prompt, vlm_model, vlm_prefill, vlm_thinking_mode, + analyze_question, analyze_system, analyze_prompt, analyze_model, + clip_model, blip_model, clip_mode, + wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape + ], + outputs=[prompt, output_image] + ) + vlm_batch = gr.State('caption') + btn_vlm_caption_batch.click( + fn=vqa.batch, + inputs=[vlm_model, vlm_system, vlm_batch_files, vlm_batch_folder, vlm_batch_str, vlm_question, vlm_prompt, vlm_save_txt, vlm_save_append, vlm_save_json, vlm_folder_recursive, vlm_prefill, vlm_thinking_mode, vlm_batch], + outputs=[prompt] + ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) + + btn_analyze_caption.click( + _js="getCaptionActiveTab", # js to insert current tab name as first argument + fn=caption_wrapper, + inputs=[dummy, image, + vlm_question, vlm_system, vlm_prompt, vlm_model, vlm_prefill, vlm_thinking_mode, + analyze_question, analyze_system, analyze_prompt, analyze_model, + clip_model, blip_model, clip_mode, + wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape + ], + outputs=[prompt, output_image] + ) + analyze_batch = gr.State('analyze') + btn_analyze_caption_batch.click( + fn=vqa.batch, + inputs=[analyze_model, analyze_system, analyze_batch_files, analyze_batch_folder, analyze_batch_str, analyze_question, analyze_prompt, analyze_save_txt, analyze_save_append, analyze_save_json,analyze_folder_recursive, vlm_prefill, vlm_thinking_mode, analyze_batch], + outputs=[prompt] + ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) + + btn_clip_caption_img.click( + fn=openclip.caption_image, + inputs=[image, clip_model, blip_model, clip_mode], + outputs=[prompt] + ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) + btn_clip_analyze_img.click( + fn=openclip.analyze_image, + inputs=[image, clip_model, blip_model], + outputs=[medium, artist, movement, trending, flavor, clip_labels_text] + ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) + + btn_wd_tag.click( + fn=tagger_tag_wrapper, + inputs=[image, wd_model, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape], + outputs=[prompt] + ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) + btn_wd_tag_batch.click( + fn=tagger_batch_wrapper, + inputs=[wd_model, wd_batch_files, wd_batch_folder, wd_batch_str, wd_save_output, wd_save_append, wd_folder_recursive, wd_general_threshold, wd_character_threshold, wd_include_rating, wd_exclude_tags, wd_max_tags, wd_sort_alpha, wd_use_spaces, wd_escape], + outputs=[prompt] + ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) + btn_clip_caption_batch.click( + fn=openclip.caption_batch, + inputs=[dummy, clip_batch_files, clip_batch_folder, clip_batch_str, clip_model, blip_model, clip_mode, clip_save_output, clip_save_append, clip_folder_recursive], + outputs=[prompt] + ).then(fn=lambda: gr.update(visible=False), inputs=[], outputs=[output_image]) # Dynamic UI updates based on selected model and task vlm_model.change(fn=update_vlm_prompts_for_model, inputs=[vlm_model], outputs=[vlm_question]) @@ -324,6 +441,9 @@ def create_ui(): # Load/Unload model buttons vlm_load_btn.click(fn=vqa.load_model, inputs=[vlm_model], outputs=[]) vlm_unload_btn.click(fn=vqa.unload_model, inputs=[], outputs=[]) + analyze_load_btn.click(fn=vqa.load_model, inputs=[vlm_model], outputs=[]) + analyze_unload_btn.click(fn=vqa.unload_model, inputs=[], outputs=[]) + def tagger_load_wrapper(model_name): from modules.caption import tagger return tagger.load_model(model_name)