feat(tagger): swap Qwen-VL for Moondream2 NL captioner

src/ensemble_tagger.py

@@ -278,17 +278,16 @@ class EnsembleTagger:
         result["pose_tags"] = pose_info.get("pose_tags", [])
         result["people_count"] = pose_info.get("people_count", 0)
 
-        # Optional Qwen3-VL natural-language caption. Lazy-loaded on first call.
-        # On HF Spaces (CPU only) this is skipped by default because the model
-        # alone is ~8-9 GB even in fp16 CPU; set WHYX_ENABLE_QWEN_VL=1 to enable.
-        if not skip_qwen and mode in ("ensemble", "qwen"):
-            try:
-                from src.qwen_vl_tagger import get_qwen_tagger
-                caption = get_qwen_tagger().describe(pil_img)
-                if caption:
-                    result["nl_caption"] = caption
-            except Exception:
-                pass
+        # Moondream2 natural-language caption (replaces Qwen-VL for HF Spaces).
+        # Runs alongside the tagger — result always carries an `nl_caption` key
+        # (empty string when Moondream is disabled or unavailable).
+        result["nl_caption"] = ""
+        try:
+            from src.moondream_tagger import get_moondream_tagger
+            nl_txt = get_moondream_tagger().caption(pil_img, length="short")
+            result["nl_caption"] = nl_txt or ""
+        except Exception:  # pragma: no cover
+            result["nl_caption"] = ""
 
         # DeepDanbooru boosts general tags in ensemble. Only runs when the optional
         # weights happen to be co-located with the app; the app keeps working when

src/moondream_tagger.py

@@ -0,0 +1,153 @@
+"""Moondream2 captioner for natural-language scene descriptions.
+
+Moondream2 is a much smaller VLM than Qwen3-VL (~3.7 GB bf16, ~2B params) and
+runs well on CPU-only HF Spaces free tier (2 vCPU / 16 GB RAM). The model card
+states the public API over the *model class itself* — caption / query / detect /
+point — so we call the high-level helpers directly instead of a generic
+"generate()" loop.
+
+By default the model is loaded lazily on first call (CPU). The runtime RAM
+footprint stays well under 8 GB when stacked alongside WD14 taggers, pose
+estimators and the SD pipeline.
+
+Set ``WHYX_DISABLE_MOONDREAM=1`` to keep Moondream off in constrained
+deployments.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import os
+import threading
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+try:
+    import torch
+    from PIL import Image, ImageOps
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+except Exception:  # pragma: no cover - keep import errors soft
+    AutoModelForCausalLM = AutoTokenizer = None
+    torch = None
+    Image = ImageOps = None
+
+_MODEL_ID = os.environ.get("WHYX_MOONDREAM_MODEL", "vikhyatk/moondream2")
+_MD_INSTANCE: "MoondreamTagger | None" = None
+
+
+def _md_enabled() -> bool:
+    return os.environ.get("WHYX_DISABLE_MOONDREAM", "0").strip().lower() not in ("1", "true", "yes", "on")
+
+
+class MoondreamTagger:
+    """One-stop wrapper around `vikhyatk/moondream2` for caption + query.
+    """
+
+    def __init__(self, model_id: str = _MODEL_ID):
+        self._model_id = model_id
+        self._model: Optional[AutoModelForCausalLM] = None
+        self._tokenizer: Optional[AutoTokenizer] = None
+        self._loaded = False
+        self._lock = threading.Lock()
+
+    def ensure_loaded(self) -> bool:
+        if self._loaded:
+            return True
+        if not _md_enabled() or AutoModelForCausalLM is None:
+            return False
+        with self._lock:
+            if self._loaded:
+                return True
+            try:
+                # bf16 halves RAM vs fp32 when a GPU is present; CPU gets fp32
+                # (still well under the 16 GB cap once the model is alive).
+                dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+                self._model = AutoModelForCausalLM.from_pretrained(
+                    self._model_id,
+                    dtype=dtype,
+                    # moondream2's anti-hallucination / EOS logic lives in its config
+                    trust_remote_code=True,
+                    # Keep all RAM on CPU; avoid any accelerate/gpu dispatch.
+                    device_map={"": "cpu"},
+                )
+                self._tokenizer = AutoTokenizer.from_pretrained(self._model_id)
+            except Exception as exc:  # pragma: no cover
+                logger.warning("Moondream2 failed to load: %s", exc)
+                return False
+            self._loaded = True
+            return True
+
+    @staticmethod
+    def _to_pil(image) -> "Image.Image":
+        # Accept PIL/numpy/raw bytes; EXIF orientation is normalised.
+        if Image is None:
+            raise RuntimeError("PIL is not available")
+        if isinstance(image, Image.Image):
+            pil = image.convert("RGB")
+        else:
+            try:
+                import numpy as np
+                if isinstance(image, np.ndarray):
+                    pil = Image.fromarray(image)
+                else:
+                    pil = Image.open(io.BytesIO(image)).convert("RGB")
+            except Exception as exc:
+                raise ValueError(f"cannot convert input to PIL image: {exc}") from exc
+        if ImageOps is not None:
+            pil = ImageOps.exif_transpose(pil)
+        return pil
+
+    def caption(self, image, length: str = "normal") -> str:
+        """Return a natural-language caption for the image.
+
+        ``length`` may be ``"short"`` (one phrase) or ``"normal"`` (one sentence).
+        """
+        if not self.ensure_loaded():
+            return ""
+        pil = self._to_pil(image)
+        try:
+            result = self._model.caption(pil, length=length)
+        except Exception as exc:  # pragma: no cover
+            logger.warning("Moondream2 caption failed: %s", exc)
+            return ""
+        # The public API returns a dict with a "caption" key. Helper text is
+        # stripped to one sentence (and one line) so it can join tags cleanly.
+        text = result.get("caption", "")
+        return " ".join(text.split())
+
+    def query(self, image, question: str) -> str:
+        """Free-form visual Q&A. E.g. "How many people are in the image?" """
+        if not self.ensure_loaded():
+            return ""
+        pil = self._to_pil(image)
+        try:
+            result = self._model.query(pil, question)
+        except Exception as exc:  # pragma: no cover
+            logger.warning("Moondream2 query failed: %s", exc)
+            return ""
+        return " ".join(result.get("answer", "").split())
+
+    def detect(self, image, thing: str) -> int:
+        """Detect instances of `thing` in `image`, return count."""
+        if not self.ensure_loaded():
+            return 0
+        pil = self._to_pil(image)
+        try:
+            result = self._model.detect(pil, thing)
+        except Exception:
+            return 0
+        return len(result.get("objects", []))
+
+
+import io
+import threading
+
+
+def get_moondream_tagger() -> MoondreamTagger:
+    """Singleton accessor shared across calls (one model per process)."""
+    global _MD_INSTANCE
+    if _MD_INSTANCE is None:
+        _MD_INSTANCE = MoondreamTagger()
+    return _MD_INSTANCE

src/qwen_vl_tagger.py

@@ -1,131 +0,0 @@
-"""Qwen3-VL based image captioner for Stable Diffusion prompts.
-
-Uses the Qwen3-VL-Instruct vision-language model (default: 4B) to produce a
-concise natural-language scene description that can be appended to the tagger
-output. Loading is opt-in because the model weighs ~4-5 GB in fp16; the class
-degrades gracefully when the DEPS or the flag are missing.
-"""
-from __future__ import annotations
-
-import os
-from typing import Optional
-
-import numpy as np
-import torch
-from PIL import Image, ImageOps
-
-_VL_DEPS_OK = True
-try:
-    from transformers import AutoProcessor
-    # Qwen3-VL needs the generic ImageTextToText class; older transformers
-    # (<4.55) only had AutoModelForVision2Seq. We try both to stay compatible.
-    try:
-        from transformers import AutoModelForImageTextToText as _AutoModel
-    except ImportError:  # pragma: no cover
-        from transformers import AutoModelForVision2Seq as _AutoModel
-except Exception:  # pragma: no cover
-    _VL_DEPS_OK = False
-    AutoProcessor = _AutoModel = None
-
-_QWEN_ENABLED = os.getenv("WHYX_ENABLE_QWEN_VL", "").strip().lower() in ("1", "true", "yes", "on")
-
-_MODEL_ID = os.environ.get("WHYX_QWEN_VL_MODEL", "Qwen/Qwen3-VL-4B-Instruct")
-_vl_instance: "QwenVLTagger | None" = None
-
-
-
-def _vl_enabled() -> bool:
-    return os.environ.get("WHYX_ENABLE_QWEN_VL", "1").strip().lower() not in ("0", "false", "no", "off")
-
-
-def _normalize_caption(raw: str) -> str:
-    """Normalize a VL caption so it joins cleanly with a Stable Diffusion prompt.
-
-    - Collapse whitespace and stray newlines.
-    - Remove leading caption markers like "This image shows..."."""
-    text = " ".join(raw.split())
-    # Strip leading meta framing if present.
-    for prefix in ("The image", "This image", "The photo", "This photo", "A scene of"):
-        if text.startswith(prefix):
-            text = text[len(prefix):].lstrip(" ,:;")
-            break
-    # Capitalize-first letter; leave the rest untouched.
-    if text:
-        text = text[0].upper() + text[1:]
-    return text
-
-
-class QwenVLTagger:
-    """Qwen3-VL wrapper exclusively for the tagger's natural-language caption.
-
-    The model is fully disabled by default (`WHYX_ENABLE_QWEN_VL=1` activates it).
-    On CPU-only environments the model would otherwise consume more than the
-    16 GB limit — hence the off default.
-    """
-
-    def __init__(self, model_id: str = _MODEL_ID):
-        self._model_id = model_id
-        self._processor: Optional[AutoProcessor] = None
-        self._model: Optional["_AutoModel"] = None
-        self._loaded = False
-
-    def ensure_loaded(self) -> bool:
-        if self._loaded:
-            return True
-        if not _VL_DEPS_OK:
-            raise RuntimeError("transformers is not installed")
-        if not _vl_enabled():
-            raise RuntimeError("Qwen VL is disabled (WHYX_ENABLE_QWEN_VL=0)")
-        self._processor = AutoProcessor.from_pretrained(self._model_id)
-        dtype = torch.float32  # FP32 on CPU — saves 4 bits per parameter over bf16.
-        self._model = _AutoModel.from_pretrained(
-            self._model_id,
-            dtype=dtype,
-        )
-        # Keep the model resident on the target device (CPU on our tier).
-        device = "cpu"
-        self._model = self._model.to(device)
-        self._loaded = True
-        return True
-
-    @staticmethod
-    def _to_pil(image) -> Image.Image:
-        img = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))
-        # EXIF transpose so rotated photo inputs are interpreted upright.
-        return ImageOps.exif_transpose(img.convert("RGB"))
-
-    def describe(self, image, prompt: str | None = None, max_new_tokens: int = 128) -> str:
-        if not self.ensure_loaded():
-            return ""
-        img = self._to_pil(image)
-        # Qwen3-VL accepts a chat-style multi-turn format; we only need one-turn
-        # caption generation. Keep the framing as user-message-only for the best
-        # instruction-following on the free tier.
-        system_prompt = prompt or (
-            "Describe this image in ONE concise sentence suitable as a Stable "
-            "Diffusion prompt (no preamble, no extra sentences)."
-        )
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": "Image:"},
-                    {"type": "image"},
-                    {"type": "text", "text": system_prompt},
-                ],
-            },
-        ]
-        text = self._processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
-        inputs = self._processor(text=[text], images=[img], return_tensors="pt")
-        with torch.inference_mode():
-            gen = self._model.generate(**inputs, max_new_tokens=max_new_tokens)
-        trimmed = gen[:, inputs["input_ids"].shape[1]:]
-        caption = self._processor.batch_decode(trimmed, skip_special_tokens=True)[0]
-        return caption.strip()
-
-
-def get_qwen_tagger() -> QwenVLTagger:
-    global _vl_instance
-    if _vl_instance is None:
-        _vl_instance = QwenVLTagger()
-    return _vl_instance