Spaces:
Running
Running
File size: 5,606 Bytes
17a1e39 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """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
|