From b625884031f6d3a760b3a4fd7683f9faf23c3e2d Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 7 Jul 2025 10:01:02 -0400 Subject: [PATCH] add gemma3n to caption/vlm and promptenhance Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 6 ++++-- modules/interrogate/vqa.py | 8 +++++++- scripts/prompt_enhance.py | 2 ++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 10488eb9f..be344276e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,8 @@ We also switched project license to **Apache-2.0** which means that SD.Next is n Feature highlights include: - **ModernUI** layout redesign which should make it more user friendly and easier to navigate - New background replacement and relightning methods using **Latent Bridge Matching** and new **PixelArt** processing filter -- Compute improvements with support for **SageAttention2++** +- New LLM/VLM models available for captioning and prompt enhance +- Compute improvements And (as always) many bugfixes and improvements to existing features! @@ -35,7 +36,8 @@ Although upgrades and existing installations are tested and should work fine! available in *img2img or control -> scripts* - Add **FLUX.1-Kontext-Dev** inpaint workflow - Support **FLUX.1** all-in-one safetensors - - Support TAESD preview and remote VAE for **HunyuanDit** + - Support **TAESD** preview and remote VAE for **HunyuanDit** + - Support for [Gemma 3n](https://huggingface.co/google/gemma-3n-E4B-it) E2B and E4B LLM/VLM models in **prompt enhance** and process **captioning** - **UI** - major update to modernui layout - redesign of the Flat UI theme diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index 4d40a8405..f569c0e12 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -25,6 +25,8 @@ vlm_models = { "CogFlorence 2.2 Large": "thwri/CogFlorence-2.2-Large", # 1.6GB "Moondream 2": "vikhyatk/moondream2", # 3.7GB "Google Gemma 3 4B": "google/gemma-3-4b-it", + "Google Gemma 3n E2B": "google/gemma-3n-E2B-it", # 1.5GB + "Google Gemma 3n E4B": "google/gemma-3n-E4B-it", # 1.5GB "Google Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB "Google PaliGemma 2 3B": "google/paligemma2-3b-pt-224", "Alibaba Qwen VL2 2B": "Qwen/Qwen2-VL-2B-Instruct", @@ -168,7 +170,11 @@ def gemma(question: str, image: Image.Image, repo: str = None, system_prompt: st if model is None or loaded != repo: shared.log.debug(f'Interrogate load: vlm="{repo}"') model = None - model = transformers.Gemma3ForConditionalGeneration.from_pretrained( + if '3n' in repo: + cls = transformers.Gemma3nForConditionalGeneration + else: + cls = transformers.Gemma3ForConditionalGeneration + model = cls.from_pretrained( repo, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index 077d8959c..874551ee0 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -35,6 +35,8 @@ class Options: models = { 'google/gemma-3-1b-it': {}, 'google/gemma-3-4b-it': {}, + 'google/gemma-3n-E2B-it': {}, + 'google/gemma-3n-E4B-it': {}, 'Qwen/Qwen3-0.6B-FP8': {}, 'Qwen/Qwen3-1.7B-FP8': {}, 'Qwen/Qwen3-4B-FP8': {},