From 30162c06022a77058110b0df01e918695b9b8409 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Wed, 15 Jul 2026 10:25:47 -0600
Subject: [PATCH] Improvements for captioner quantization to speed it up. Block
 compile on captioners.

---
 .../captioner/BaseCaptioner.py                | 43 ++++++++++++++++++-
 .../captioner/Qwen3VLCaptioner.py             | 22 ++++++++--
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py
index 032b5f8e..6100fbbd 100644
--- a/extensions_built_in/captioner/BaseCaptioner.py
+++ b/extensions_built_in/captioner/BaseCaptioner.py
@@ -208,14 +208,53 @@ class BaseCaptioner(BaseExtensionProcess):
             torch._dynamo.config.suppress_errors = True
             for model in [self.model, self.model2]:
                 if model is not None and isinstance(model, torch.nn.Module):
-                    # dynamic=True avoids recompiling for every new image/token shape
-                    model.compile(dynamic=True)
+                    # compile per transformer block instead of the whole model:
+                    # small graphs compile far faster and identical blocks hit
+                    # the inductor cache, vs many minutes tracing one huge graph
+                    compiled_blocks = self._compile_blocks(model)
+                    if compiled_blocks == 0:
+                        # no repeated block lists found; compile the whole model
+                        # dynamic=True avoids recompiling for every new image/token shape
+                        model.compile(dynamic=True)
             print(
                 "[AITK] Model compilation enabled. The first few items will be slow while the model compiles."
             )
         except Exception as e:
             print(f"[AITK] Failed to compile model, continuing without compile: {e}")
 
+    def _compile_blocks(self, model: torch.nn.Module) -> int:
+        """Compile the repeated transformer blocks individually, leaving one-off
+        modules (embeddings, mergers, lm_head) eager. Returns the number of
+        blocks compiled."""
+        # candidate lists: ModuleLists of >= 2 blocks that all share one class
+        # and have submodules of their own (i.e. real transformer blocks, not
+        # lists of leaf layers)
+        candidates = []
+        for name, module in model.named_modules():
+            if not isinstance(module, torch.nn.ModuleList) or len(module) < 2:
+                continue
+            classes = {type(b) for b in module}
+            if len(classes) != 1:
+                continue
+            if next(module[0].children(), None) is None:
+                continue
+            candidates.append(name)
+        # skip lists nested inside another candidate list
+        candidates = [
+            name
+            for name in candidates
+            if not any(
+                name != other and name.startswith(other + ".") for other in candidates
+            )
+        ]
+        count = 0
+        for name in candidates:
+            block_list = model.get_submodule(name)
+            for i, block in enumerate(block_list):
+                block_list[i] = torch.compile(block, dynamic=True)
+                count += 1
+        return count
+
     def start_stop_watcher(self, interval_sec: float = 5.0):
         """
         Start a daemon thread that periodically checks should_stop()
diff --git a/extensions_built_in/captioner/Qwen3VLCaptioner.py b/extensions_built_in/captioner/Qwen3VLCaptioner.py
index 88bf41b0..3c8a5c42 100644
--- a/extensions_built_in/captioner/Qwen3VLCaptioner.py
+++ b/extensions_built_in/captioner/Qwen3VLCaptioner.py
@@ -28,10 +28,10 @@ def patch_qwen_vl_patch_embed(model):
     patched = 0
     for module in model.modules():
         proj = getattr(module, "proj", None)
-        if (
-            isinstance(proj, torch.nn.Conv3d)
-            and tuple(proj.kernel_size) == tuple(proj.stride)
+        if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple(
+            proj.stride
         ):
+
             def fast_forward(hidden_states, _proj=proj):
                 w = _proj.weight.reshape(_proj.weight.shape[0], -1)
                 x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
@@ -41,6 +41,7 @@ def patch_qwen_vl_patch_embed(model):
             patched += 1
     return patched
 
+
 # transformers.logging.set_verbosity_error()
 warnings.filterwarnings("ignore")
 logging.disable(logging.WARNING)
@@ -84,7 +85,20 @@ class Qwen3VLCaptioner(BaseCaptioner):
             self.model.to(self.device_torch)
         if self.caption_config.quantize:
             self.print_and_status_update("Quantizing Qwen3VL model")
-            quantize(self.model, weights=get_qtype(self.caption_config.qtype))
+            # in low vram mode the model stays on cpu; quantize each layer on the
+            # gpu and move it back so the math is fast without holding the whole
+            # model in vram
+            # lm_head is huge (vocab x hidden) and quality-critical; quantizing it
+            # needs a ~4x transient allocation that can OOM, so keep it in full
+            # precision
+            quantize(
+                self.model,
+                weights=get_qtype(self.caption_config.qtype),
+                exclude=["lm_head", "*.lm_head"],
+                quantize_device=self.device_torch
+                if self.caption_config.low_vram
+                else None,
+            )
             freeze(self.model)
             flush()
         self.processor = AutoProcessor.from_pretrained(