From 80bb331169fcf36f2c3ba0b18e4379d361a4a690 Mon Sep 17 00:00:00 2001
From: CalamitousFelicitousness <iowasovereign@gmail.com>
Date: Sun, 26 Oct 2025 06:01:33 +0000
Subject: [PATCH] Prompt enhance resizing and Qwen VL fix

---
 modules/interrogate/vqa.py |  6 +++---
 scripts/prompt_enhance.py  | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py
index 0275773ce..750439f0c 100644
--- a/modules/interrogate/vqa.py
+++ b/modules/interrogate/vqa.py
@@ -187,11 +187,11 @@ def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
         model = None
-        if 'Qwen3' in repo:
+        if 'Qwen3-VL' in repo or 'Qwen3VL' in repo:
             cls_name = transformers.Qwen3VLForConditionalGeneration
-        elif '2.5' in repo:
+        elif 'Qwen2.5-VL' in repo or 'Qwen2_5_VL' in repo:
             cls_name = transformers.Qwen2_5_VLForConditionalGeneration
-        else:
+        elif 'Qwen2-VL' in repo or 'Qwen2VL' in repo:
             cls_name = transformers.Qwen2VLForConditionalGeneration
         model = cls_name.from_pretrained(
             repo,
diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py
index 87396136f..ec5d0b3e6 100644
--- a/scripts/prompt_enhance.py
+++ b/scripts/prompt_enhance.py
@@ -168,7 +168,19 @@ class Script(scripts_manager.Script):
             load_args = { 'pretrained_model_name_or_path': model_repo if not gguf_args else model_gguf }
             if model_subfolder:
                 load_args['subfolder'] = model_subfolder # Comma was incorrect here
-            self.llm = transformers.AutoModelForCausalLM.from_pretrained(
+
+            # Determine model class based on model type
+            if 'Qwen3-VL' in model_repo or 'Qwen3VL' in model_repo:
+                model_cls = transformers.Qwen3VLForConditionalGeneration
+                # Use Qwen3-VL's optimized default attention
+            elif 'Qwen2.5-VL' in model_repo or 'Qwen2_5_VL' in model_repo:
+                model_cls = transformers.Qwen2_5_VLForConditionalGeneration
+            elif 'Qwen2-VL' in model_repo or 'Qwen2VL' in model_repo:
+                model_cls = transformers.Qwen2VLForConditionalGeneration
+            else:
+                model_cls = transformers.AutoModelForCausalLM
+
+            self.llm = model_cls.from_pretrained(
                 **load_args,
                 trust_remote_code=True,
                 torch_dtype=devices.dtype,
@@ -302,6 +314,25 @@ class Script(scripts_manager.Script):
         except Exception:
             current_image = None
 
+        # Resize large images to match VQA performance (vision models are sensitive to resolution)
+        # Create a copy to avoid modifying the original image used by img2img
+        if current_image is not None and isinstance(current_image, Image.Image):
+            original_size = (current_image.width, current_image.height)
+            needs_resize = current_image.width > 768 or current_image.height > 768
+            needs_rgb = current_image.mode != 'RGB'
+
+            if needs_resize or needs_rgb:
+                # Copy the image before any modifications to preserve the original
+                current_image = current_image.copy()
+
+                if needs_resize:
+                    current_image.thumbnail((768, 768), Image.Resampling.LANCZOS)
+                    debug_log(f'Prompt enhance: Resized image from {original_size} to {(current_image.width, current_image.height)}')
+
+                if needs_rgb:
+                    current_image = current_image.convert('RGB')
+                    debug_log(f'Prompt enhance: Converted image to RGB mode')
+
         has_system = system is not None and len(system) > 4
         mode = 'custom' if has_system else ''