inference-models 0.18.3__py3-none-any.whl

inference_models/models/moondream2/moondream2_hf.py

@@ -0,0 +1,281 @@
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models import Detections
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat, ImageDimensions
+from inference_models.errors import ModelRuntimeError
+from inference_models.models.common.model_packages import get_model_package_contents
+from inference_models.models.common.roboflow.pre_processing import images_to_pillow
+from inference_models.utils.imports import import_class_from_file
+
+
+@dataclass
+class EncodedImage:
+    moondream_encoded_image: Any
+    image_dimensions: ImageDimensions
+
+
+@dataclass
+class Points:
+    xy: torch.Tensor
+    confidence: torch.Tensor
+    class_id: torch.Tensor
+
+
+class MoonDream2HF:
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        **kwargs,
+    ) -> "MoonDream2HF":
+        if torch.mps.is_available():
+            raise ModelRuntimeError(
+                message=f"This model cannot run on Apple device with MPS unit - original implementation contains bug "
+                f"preventing proper allocation of tensors which causes runtime error. Run this model on the "
+                f"machine with Nvidia GPU or x86 CPU.",
+                help_url="https://todo",
+            )
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=["hf_moondream.py"],
+        )
+        model_class = import_class_from_file(
+            file_path=model_package_content["hf_moondream.py"],
+            class_name="HfMoondream",
+        )
+        model = model_class.from_pretrained(model_name_or_path).to(device)
+        return cls(model=model, device=device)
+
+    def __init__(self, model, device: torch.device):
+        self._model = model
+        self._device = device
+
+    def detect(
+        self,
+        images: Union[
+            EncodedImage,
+            List[EncodedImage],
+            torch.Tensor,
+            List[torch.Tensor],
+            np.ndarray,
+            List[np.ndarray],
+        ],
+        classes: List[str],
+        max_tokens: int = 700,
+        input_color_format: Optional[ColorFormat] = None,
+    ) -> List[Detections]:
+        encoded_images = self.encode_images(
+            images=images, input_color_format=input_color_format
+        )
+        results = []
+        for encoded_image in encoded_images:
+            image_detections = []
+            for class_id, class_name in enumerate(classes):
+                class_detections = self._model.detect(
+                    image=encoded_image.moondream_encoded_image,
+                    object=class_name,
+                    settings={"max_tokens": max_tokens},
+                )["objects"]
+                image_detections.append((class_id, class_detections))
+            image_results = post_process_detections(
+                raw_detections=image_detections,
+                image_dimensions=encoded_image.image_dimensions,
+                device=self._device,
+            )
+            results.append(image_results)
+        return results
+
+    def caption(
+        self,
+        images: Union[
+            EncodedImage,
+            List[EncodedImage],
+            torch.Tensor,
+            List[torch.Tensor],
+            np.ndarray,
+            List[np.ndarray],
+        ],
+        length: Literal["normal", "short", "long"] = "normal",
+        max_tokens: int = 700,
+        input_color_format: Optional[ColorFormat] = None,
+    ) -> List[str]:
+        encoded_images = self.encode_images(
+            images=images, input_color_format=input_color_format
+        )
+        results = []
+        for encoded_image in encoded_images:
+            result = self._model.caption(
+                image=encoded_image.moondream_encoded_image,
+                length=length,
+                settings={"max_tokens": max_tokens},
+            )
+            results.append(result["caption"].strip())
+        return results
+
+    def query(
+        self,
+        images: Union[
+            EncodedImage,
+            List[EncodedImage],
+            torch.Tensor,
+            List[torch.Tensor],
+            np.ndarray,
+            List[np.ndarray],
+        ],
+        question: str,
+        max_tokens: int = 700,
+        input_color_format: Optional[ColorFormat] = None,
+    ) -> List[str]:
+        encoded_images = self.encode_images(
+            images=images, input_color_format=input_color_format
+        )
+        results = []
+        for encoded_image in encoded_images:
+            result = self._model.query(
+                image=encoded_image.moondream_encoded_image,
+                question=question,
+                settings={"max_tokens": max_tokens},
+            )
+            results.append(result["answer"].strip())
+        return results
+
+    def point(
+        self,
+        images: Union[
+            EncodedImage,
+            List[EncodedImage],
+            torch.Tensor,
+            List[torch.Tensor],
+            np.ndarray,
+            List[np.ndarray],
+        ],
+        classes: List[str],
+        max_tokens: int = 700,
+        input_color_format: Optional[ColorFormat] = None,
+    ) -> List[Points]:
+        encoded_images = self.encode_images(
+            images=images, input_color_format=input_color_format
+        )
+        results = []
+        for encoded_image in encoded_images:
+            image_points = []
+            for class_id, class_name in enumerate(classes):
+                class_points = self._model.point(
+                    image=encoded_image.moondream_encoded_image,
+                    object=class_name,
+                    settings={"max_tokens": max_tokens},
+                )["points"]
+                image_points.append((class_id, class_points))
+            image_results = post_process_points(
+                raw_points=image_points,
+                image_dimensions=encoded_image.image_dimensions,
+                device=self._device,
+            )
+            results.append(image_results)
+        return results
+
+    def encode_images(
+        self,
+        images: Union[
+            EncodedImage,
+            List[EncodedImage],
+            torch.Tensor,
+            List[torch.Tensor],
+            np.ndarray,
+            List[np.ndarray],
+        ],
+        input_color_format: Optional[ColorFormat] = None,
+    ) -> List[EncodedImage]:
+        if are_images_encoded(images=images):
+            if not isinstance(images, list):
+                return [images]
+            return images
+        pillow_images, images_dimensions = images_to_pillow(
+            images=images,
+            input_color_format=input_color_format,
+            model_color_format="rgb",
+        )
+        result = []
+        for image, image_dimensions in zip(pillow_images, images_dimensions):
+            moondream_encoded = self._model.encode_image(image)
+            result.append(
+                EncodedImage(
+                    moondream_encoded_image=moondream_encoded,
+                    image_dimensions=image_dimensions,
+                )
+            )
+        return result
+
+
+def are_images_encoded(
+    images: Union[
+        EncodedImage,
+        List[EncodedImage],
+        torch.Tensor,
+        List[torch.Tensor],
+        np.ndarray,
+        List[np.ndarray],
+    ],
+) -> bool:
+    if isinstance(images, list):
+        if not len(images):
+            raise ModelRuntimeError(
+                message="Detected empty input to the model", help_url="https://todo"
+            )
+        return isinstance(images[0], EncodedImage)
+    return isinstance(images, EncodedImage)
+
+
+def post_process_detections(
+    raw_detections: List[Tuple[int, List[dict]]],
+    image_dimensions: ImageDimensions,
+    device: torch.device,
+) -> Detections:
+    xyxy, confidence, class_id = [], [], []
+    for detection_class_id, raw_class_detections in raw_detections:
+        for raw_detection in raw_class_detections:
+            xyxy.append(
+                [
+                    raw_detection["x_min"] * image_dimensions.width,
+                    raw_detection["y_min"] * image_dimensions.height,
+                    raw_detection["x_max"] * image_dimensions.width,
+                    raw_detection["y_max"] * image_dimensions.height,
+                ]
+            )
+            class_id.append(detection_class_id)
+            confidence.append(1.0)
+    return Detections(
+        xyxy=torch.tensor(xyxy, device=device).round().int(),
+        class_id=torch.tensor(class_id, device=device).int(),
+        confidence=torch.tensor(confidence, device=device),
+    )
+
+
+def post_process_points(
+    raw_points: List[Tuple[int, List[dict]]],
+    image_dimensions: ImageDimensions,
+    device: torch.device,
+) -> Points:
+    xy, confidence, class_id = [], [], []
+    for point_class_id, raw_class_points in raw_points:
+        for raw_point in raw_class_points:
+            xy.append(
+                [
+                    raw_point["x"] * image_dimensions.width,
+                    raw_point["y"] * image_dimensions.height,
+                ]
+            )
+            class_id.append(point_class_id)
+            confidence.append(1.0)
+    return Points(
+        xy=torch.tensor(xy, device=device).round().int(),
+        class_id=torch.tensor(class_id, device=device).int(),
+        confidence=torch.tensor(confidence, device=device),
+    )

inference_models/models/owlv2/cache.py

@@ -0,0 +1,182 @@
+import hashlib
+import json
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from threading import Lock
+from typing import Dict, List, Optional
+
+import torch
+
+from inference_models.errors import EnvironmentConfigurationError
+from inference_models.models.owlv2.entities import (
+    ImageEmbeddings,
+    LazyReferenceExample,
+    ReferenceExamplesClassEmbeddings,
+)
+
+
+class OwlV2ClassEmbeddingsCache(ABC):
+
+    @abstractmethod
+    def retrieve_embeddings(
+        self, key: str
+    ) -> Optional[Dict[str, ReferenceExamplesClassEmbeddings]]:
+        pass
+
+    @abstractmethod
+    def save_embeddings(
+        self, key: str, embeddings: Dict[str, ReferenceExamplesClassEmbeddings]
+    ) -> None:
+        pass
+
+
+class OwlV2ClassEmbeddingsCacheNullObject(OwlV2ClassEmbeddingsCache):
+
+    def retrieve_embeddings(
+        self, key: str
+    ) -> Optional[Dict[str, ReferenceExamplesClassEmbeddings]]:
+        return None
+
+    def save_embeddings(
+        self, key: str, embeddings: Dict[str, ReferenceExamplesClassEmbeddings]
+    ) -> None:
+        pass
+
+
+class InMemoryOwlV2ClassEmbeddingsCache(OwlV2ClassEmbeddingsCache):
+
+    @classmethod
+    def init(
+        cls, size_limit: Optional[int], send_to_cpu: bool = True
+    ) -> "InMemoryOwlV2ClassEmbeddingsCache":
+        return cls(
+            state=OrderedDict(),
+            size_limit=size_limit,
+            send_to_cpu=send_to_cpu,
+        )
+
+    def __init__(
+        self,
+        state: OrderedDict,
+        size_limit: Optional[int],
+        send_to_cpu: bool = True,
+    ):
+        self._state = state
+        self._size_limit = size_limit
+        self._send_to_cpu = send_to_cpu
+        self._state_lock = Lock()
+
+    def retrieve_embeddings(
+        self, key: str
+    ) -> Optional[Dict[str, ReferenceExamplesClassEmbeddings]]:
+        return self._state.get(key)
+
+    def save_embeddings(
+        self, key: str, embeddings: Dict[str, ReferenceExamplesClassEmbeddings]
+    ) -> None:
+        with self._state_lock:
+            if key in self._state:
+                return None
+            self._ensure_cache_has_capacity()
+            if self._send_to_cpu:
+                embeddings = {
+                    k: v.to(device=torch.device("cpu")) for k, v in embeddings.items()
+                }
+            self._state[key] = embeddings
+
+    def _ensure_cache_has_capacity(self) -> None:
+        if self._size_limit < 1:
+            raise EnvironmentConfigurationError(
+                message=f"In memory cache size for OWLv2 embeddings was set to invalid value. "
+                f"If you are running inference locally - adjust settings of your deployment. If you see this "
+                f"error running on Roboflow platform - contact us to get help.",
+                help_url="https://todo",
+            )
+        if self._size_limit is None or self._size_limit < 1:
+            return None
+        while len(self._state) > self._size_limit:
+            _ = self._state.popitem(last=False)
+
+
+class OwlV2ImageEmbeddingsCache(ABC):
+
+    @abstractmethod
+    def retrieve_embeddings(self, key: str) -> Optional[ImageEmbeddings]:
+        pass
+
+    @abstractmethod
+    def save_embeddings(self, embeddings: ImageEmbeddings) -> None:
+        pass
+
+
+class OwlV2ImageEmbeddingsCacheNullObject(OwlV2ImageEmbeddingsCache):
+
+    def retrieve_embeddings(self, key: str) -> Optional[ImageEmbeddings]:
+        return None
+
+    def save_embeddings(self, embeddings: ImageEmbeddings) -> None:
+        pass
+
+
+class InMemoryOwlV2ImageEmbeddingsCache(OwlV2ImageEmbeddingsCache):
+
+    @classmethod
+    def init(
+        cls, size_limit: Optional[int], send_to_cpu: bool = True
+    ) -> "InMemoryOwlV2ImageEmbeddingsCache":
+        return cls(
+            state=OrderedDict(),
+            size_limit=size_limit,
+            send_to_cpu=send_to_cpu,
+        )
+
+    def __init__(
+        self,
+        state: OrderedDict,
+        size_limit: Optional[int],
+        send_to_cpu: bool = True,
+    ):
+        self._state = state
+        self._size_limit = size_limit
+        self._send_to_cpu = send_to_cpu
+        self._state_lock = Lock()
+
+    def retrieve_embeddings(self, key: str) -> Optional[ImageEmbeddings]:
+        return self._state.get(key)
+
+    def save_embeddings(self, embeddings: ImageEmbeddings) -> None:
+        with self._state_lock:
+            if embeddings.image_hash in self._state:
+                return None
+            self._ensure_cache_has_capacity()
+            if self._send_to_cpu:
+                embeddings = embeddings.to(device=torch.device("cpu"))
+            self._state[embeddings.image_hash] = embeddings
+
+    def _ensure_cache_has_capacity(self) -> None:
+        if self._size_limit < 1:
+            raise EnvironmentConfigurationError(
+                message=f"In memory cache size for OWLv2 embeddings was set to invalid value. "
+                f"If you are running inference locally - adjust settings of your deployment. If you see this "
+                f"error running on Roboflow platform - contact us to get help.",
+                help_url="https://todo",
+            )
+        if self._size_limit is None or self._size_limit < 1:
+            return None
+        while len(self._state) > self._size_limit:
+            _ = self._state.popitem(last=False)
+
+
+def hash_reference_examples(reference_examples: List[LazyReferenceExample]) -> str:
+    result = hashlib.sha1()
+    for example in reference_examples:
+        image_hash = example.image.get_hash()
+        result.update(image_hash.encode())
+        bboxes_hash_base = "---".join(
+            [
+                json.dumps(box.model_dump(), sort_keys=True, separators=(",", ":"))
+                for box in example.boxes
+            ]
+        )
+        result.update(bboxes_hash_base.encode())
+    return result.hexdigest()

inference_models/models/owlv2/entities.py

@@ -0,0 +1,112 @@
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from pydantic import BaseModel, ConfigDict, Field
+
+from inference_models.models.owlv2.reference_dataset import LazyImageWrapper
+
+POSITIVE_EXAMPLE = "positive"
+NEGATIVE_EXAMPLE = "negative"
+
+
+class ReferenceBoundingBox(BaseModel):
+    x: Union[float, int]
+    y: Union[float, int]
+    w: Union[float, int]
+    h: Union[float, int]
+    cls: str
+    negative: bool = Field(default=False)
+    absolute: bool = Field(default=True)
+
+    def to_tuple(
+        self, image_wh: Optional[Tuple[int, int]] = None
+    ) -> Tuple[
+        Union[int, float], Union[int, float], Union[int, float], Union[int, float]
+    ]:
+        if image_wh is None or self.absolute is False:
+            return self.x, self.y, self.w, self.h
+        max_dim = max(image_wh)
+        return (
+            self.x / max_dim,
+            self.y / max_dim,
+            self.w / max_dim,
+            self.h / max_dim,
+        )
+
+
+class ReferenceExample(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    image: Union[np.ndarray, torch.Tensor, str, bytes]
+    boxes: List[ReferenceBoundingBox]
+
+
+@dataclass(frozen=True)
+class LazyReferenceExample:
+    image: LazyImageWrapper
+    boxes: List[ReferenceBoundingBox]
+
+
+@dataclass(frozen=True)
+class ImageEmbeddings:
+    image_hash: str
+    objectness: torch.Tensor
+    boxes: torch.Tensor
+    image_class_embeddings: torch.Tensor
+    logit_shift: torch.Tensor
+    logit_scale: torch.Tensor
+    image_size_wh: Tuple[int, int]
+
+    def to(self, device: torch.device) -> "ImageEmbeddings":
+        return ImageEmbeddings(
+            image_hash=self.image_hash,
+            objectness=self.objectness.to(device=device),
+            boxes=self.boxes.to(device=device),
+            image_class_embeddings=self.image_class_embeddings.to(device=device),
+            logit_shift=self.logit_shift.to(device=device),
+            logit_scale=self.logit_scale.to(device=device),
+            image_size_wh=self.image_size_wh,
+        )
+
+
+@dataclass(frozen=True)
+class ReferenceExamplesClassEmbeddings:
+    positive: Optional[torch.Tensor]
+    negative: Optional[torch.Tensor]
+
+    def to(self, device: torch.device) -> "ReferenceExamplesClassEmbeddings":
+        return ReferenceExamplesClassEmbeddings(
+            positive=(
+                self.positive.to(device=device) if self.positive is not None else None
+            ),
+            negative=(
+                self.negative.to(device=device) if self.negative is not None else None
+            ),
+        )
+
+
+@dataclass(frozen=True)
+class ReferenceExamplesEmbeddings:
+    class_embeddings: Dict[str, ReferenceExamplesClassEmbeddings]
+    image_embeddings: Optional[Dict[str, ImageEmbeddings]]
+
+    @classmethod
+    def from_class_embeddings_dict(
+        cls,
+        class_embeddings: Dict[str, Dict[str, torch.Tensor]],
+        device: torch.device,
+    ) -> "ReferenceExamplesEmbeddings":
+        result = {}
+        for class_name, examples_class_embeddings in class_embeddings.items():
+            positive, negative = None, None
+            if POSITIVE_EXAMPLE in examples_class_embeddings:
+                positive = examples_class_embeddings[POSITIVE_EXAMPLE]
+            if NEGATIVE_EXAMPLE in examples_class_embeddings:
+                negative = examples_class_embeddings[NEGATIVE_EXAMPLE]
+            single_class_embeddings = ReferenceExamplesClassEmbeddings(
+                positive=positive, negative=negative
+            ).to(device=device)
+            result[class_name] = single_class_embeddings
+        return cls(class_embeddings=result, image_embeddings=None)