moondream 0.0.1__py3-none-any.whl

moondream/__init__.py

@@ -0,0 +1 @@
+from .vl import VL

moondream/preprocess.py

@@ -0,0 +1,62 @@
+import numpy as np
+from typing import List, Tuple, Union
+from PIL import Image
+
+
+def im_resize(
+    image: Image.Image,
+    size: Tuple[int, int],
+    resample: int = Image.Resampling.BICUBIC,
+) -> Image.Image:
+    return image.resize(size, resample=resample)
+
+
+def normalize(
+    image: np.ndarray,
+    mean: List[float] = [0.5, 0.5, 0.5],
+    std: List[float] = [0.5, 0.5, 0.5],
+) -> np.ndarray:
+    """
+    Normalize an image array.
+    """
+    return (image - np.array(mean)) / np.array(std)
+
+
+def create_patches(image: Image.Image, image_patch_size=378) -> np.ndarray:
+    """
+    Split the given image into a variable number of patches depending upon its
+    resolution.
+    """
+    # Start off with the global patch.
+    patches = [im_resize(image, (image_patch_size, image_patch_size))]
+
+    # Find the closest resolution template.
+    res_templates = [(1, 2), (2, 1), (2, 2)]
+    im_width, im_height = image.size
+    max_dim = max(im_width, im_height)
+    if max_dim < image_patch_size * 1.4:
+        # If the image is already small, we just do a single patch that is a
+        # duplicate of the global patch. This creates a small amount of
+        # redundant computation now, but it is simpler and future-proofs us
+        # if/when we condition the vision encoder on the patch type.
+        patches.append(patches[0])
+    else:
+        aspect_ratio = im_width / im_height
+        res_template = min(
+            res_templates, key=lambda size: abs((size[1] / size[0]) - aspect_ratio)
+        )
+        # TODO: Actually implement patching... just going to put in the global
+        # patch for now to make progress on other aspects.
+        patches.append(patches[0])
+
+    return np.stack(
+        [
+            normalize(
+                (np.array(patch_img) / 255.0),
+                mean=[0.5, 0.5, 0.5],
+                std=[0.5, 0.5, 0.5],
+            ).transpose(2, 0, 1)
+            for patch_img in patches
+        ],
+        dtype=np.float16,
+    )

moondream/vl.py

@@ -0,0 +1,271 @@
+import onnx
+import onnxruntime as ort
+import numpy as np
+import os
+import tarfile
+import json
+
+from typing import Generator, List, Union, Optional, Dict, TypedDict, Any
+from PIL import Image
+from dataclasses import dataclass
+from io import BytesIO
+from tokenizers import Tokenizer
+
+from .preprocess import create_patches
+
+
+@dataclass
+class EncodedImage:
+    kv_caches: List[np.ndarray]
+
+
+SamplingSettings = TypedDict(
+    "SamplingSettings",
+    {"max_tokens": int},
+    total=False,
+)
+
+CaptionOutput = TypedDict(
+    "CaptionOutput", {"caption": Union[str, Generator[str, None, None]]}
+)
+QueryOutput = TypedDict(
+    "QueryOutput", {"answer": Union[str, Generator[str, None, None]]}
+)
+
+DEFAULT_MAX_TOKENS = 1024
+LATEST_SUPPORTED_VERSION = 0
+
+
+class Region:
+    pass
+
+
+class VL:
+    def __init__(self, model_path: Optional[str], ort_settings: Dict[str, Any] = {}):
+        """
+        Initialize the Moondream VL (Vision Language) model.
+
+        Args:
+            model_path (str): The path to the model file.
+
+        Returns:
+            None
+        """
+
+        if model_path is None or not os.path.isfile(model_path):
+            raise ValueError("Model path is invalid or file does not exist.")
+
+        if not tarfile.is_tarfile(model_path):
+            raise ValueError(
+                "Model format not recognized. You may need to upgrade the moondream package."
+            )
+
+        self.text_decoders = []
+
+        with tarfile.open(model_path, "r:*") as tar:
+            for member in tar.getmembers():
+                name = member.name.split("/")[-1]
+
+                f = tar.extractfile(member)
+                if f is not None:
+                    contents = f.read()
+                else:
+                    continue
+
+                if name == "vision_encoder.onnx":
+                    self.vision_encoder = ort.InferenceSession(contents, **ort_settings)
+                elif name == "vision_projection.onnx":
+                    self.vision_projection = ort.InferenceSession(
+                        contents, **ort_settings
+                    )
+                elif name == "text_encoder.onnx":
+                    self.text_encoder = ort.InferenceSession(contents, **ort_settings)
+                elif "text_decoder" in name and name.endswith(".onnx"):
+                    self.text_decoders.append(
+                        ort.InferenceSession(contents, **ort_settings)
+                    )
+                elif name == "tokenizer.json":
+                    self.tokenizer = Tokenizer.from_buffer(contents)
+                elif name == "initial_kv_caches.npy":
+                    self.initial_kv_caches = [x for x in np.load(BytesIO(contents))]
+                elif name == "config.json":
+                    self.config = json.loads(contents)
+
+        assert self.vision_encoder is not None
+        assert self.vision_projection is not None
+        assert self.text_encoder is not None
+        assert len(self.text_decoders) > 0
+        assert self.tokenizer is not None
+        assert self.initial_kv_caches is not None
+        assert self.config is not None
+
+        if type(self.config) != dict or "model_version" not in self.config:
+            raise ValueError("Model format not recognized.")
+        if self.config["model_version"] > LATEST_SUPPORTED_VERSION:
+            raise ValueError(
+                "Model version not supported. You may need to upgrade the moondream package."
+            )
+
+        self.special_tokens = self.config["special_tokens"]
+        self.templates = self.config["templates"]
+
+    def encode_image(self, image: Union[Image.Image, EncodedImage]) -> EncodedImage:
+        """
+        Preprocess the image by running it through the model.
+
+        This method is useful if the user wants to make multiple queries with the same image.
+        The output is not guaranteed to be backward-compatible across version updates,
+        and should not be persisted out of band.
+
+        Args:
+            image (Image.Image): The input image to be encoded.
+
+        Returns:
+            The encoded representation of the image.
+        """
+        if type(image) == EncodedImage:
+            return image
+
+        image_patches = create_patches(image)  # type: ignore
+
+        patch_emb = self.vision_encoder.run(None, {"input": image_patches})[0]
+        patch_emb = np.concatenate([patch_emb[0], patch_emb[1]], axis=-1)
+        patch_emb = np.expand_dims(patch_emb, axis=0)
+        (inputs_embeds,) = self.vision_projection.run(None, {"input": patch_emb})
+
+        kv_caches = self.initial_kv_caches
+        for i, decoder in enumerate(self.text_decoders):
+            inputs_embeds, kv_caches[i] = decoder.run(
+                None,
+                {
+                    "inputs_embeds": inputs_embeds,
+                    "kv_cache": kv_caches[i],
+                },
+            )
+        return EncodedImage(kv_caches=kv_caches)
+
+    def _generate(
+        self, hidden: np.ndarray, encoded_image: EncodedImage, max_tokens: int
+    ) -> Generator[str, None, None]:
+        kv_caches = {
+            i: encoded_image.kv_caches[i] for i in range(len(self.text_decoders))
+        }
+
+        generated_tokens = 0
+        while generated_tokens < max_tokens:
+            for i, decoder in enumerate(self.text_decoders):
+                hidden, kv_caches[i] = decoder.run(
+                    None, {"inputs_embeds": hidden, "kv_cache": kv_caches[i]}
+                )
+
+            next_token = np.argmax(hidden, axis=-1)[0]
+            if next_token == self.special_tokens["eos"]:
+                break
+
+            yield self.tokenizer.decode([next_token])
+            generated_tokens += 1
+            (hidden,) = self.text_encoder.run(None, {"input_ids": [[next_token]]})
+
+    def caption(
+        self,
+        image: Union[Image.Image, EncodedImage],
+        length: str = "normal",
+        stream: bool = False,
+        settings: Optional[SamplingSettings] = None,
+    ) -> CaptionOutput:
+        """
+        Generate a caption for the input image.
+
+        Args:
+            image (Union[Image.Image, EncodedImage]): The input image to be captioned.
+            settings (Optional[SamplingSettings]): Optional settings for the caption generation.
+                If not provided, default settings will be used.
+
+        Returns:
+            str: The caption for the input image.
+        """
+        if "caption" not in self.templates:
+            raise ValueError("Model does not support captioning.")
+        if length not in self.templates["caption"]:
+            raise ValueError(f"Model does not support caption length '{length}'.")
+
+        (input_embeds,) = self.text_encoder.run(
+            None, {"input_ids": [self.templates["caption"][length]]}
+        )
+        if settings is None:
+            settings = {}
+        max_tokens = settings.get("max_tokens", DEFAULT_MAX_TOKENS)
+
+        encoded_image = self.encode_image(image)
+
+        def generator():
+            for t in self._generate(input_embeds, encoded_image, max_tokens):
+                yield t
+
+        if stream:
+            return {"caption": generator()}
+        else:
+            out = ""
+            for t in generator():
+                out += t
+            return {"caption": out}
+
+    def query(
+        self,
+        image: Union[Image.Image, EncodedImage],
+        question: str,
+        stream: bool = False,
+        settings: Optional[SamplingSettings] = None,
+    ) -> QueryOutput:
+        """
+        Generate an answer to the input question about the input image.
+
+        Args:
+            image (Union[Image.Image, EncodedImage]): The input image to be queried.
+            question (str): The question to be answered.
+
+        Returns:
+            str: The answer to the input question about the input image.
+        """
+        if "query" not in self.templates:
+            raise ValueError("Model does not support querying.")
+
+        question_toks = (
+            self.templates["query"]["prefix"]
+            + self.tokenizer.encode(question).ids
+            + self.templates["query"]["suffix"]
+        )
+
+        (input_embeds,) = self.text_encoder.run(None, {"input_ids": [question_toks]})
+        if settings is None:
+            settings = {}
+        max_tokens = settings.get("max_tokens", DEFAULT_MAX_TOKENS)
+
+        encoded_image = self.encode_image(image)
+
+        def generator():
+            for t in self._generate(input_embeds, encoded_image, max_tokens):
+                yield t
+
+        if stream:
+            return {"answer": generator()}
+        else:
+            out = ""
+            for t in generator():
+                out += t
+            return {"answer": out}
+
+    def detect(
+        self, image: Union[Image.Image, EncodedImage], object: str
+    ) -> List[Region]:
+        """
+        Detect and localize the specified object in the input image.
+
+        Args:
+            image (Union[Image.Image, EncodedImage]): The input image to be analyzed.
+            object (str): The object to be detected in the image.
+
+        Returns:
+            List[Region]: A list of Region objects representing the detected instances of the specified object.
+        """
+        return []

moondream-0.0.1.dist-info/METADATA

@@ -0,0 +1,82 @@
+Metadata-Version: 2.1
+Name: moondream
+Version: 0.0.1
+Summary: Python client library for moondream
+Author: vik
+Author-email: vik@moondream.ai
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: numpy (>=2.1.2,<3.0.0)
+Requires-Dist: onnx (>=1.17.0,<2.0.0)
+Requires-Dist: onnxruntime (>=1.19.2,<2.0.0)
+Requires-Dist: pillow (>=10.4.0,<11.0.0)
+Requires-Dist: tokenizers (>=0.20.1,<0.21.0)
+Description-Content-Type: text/markdown
+
+# Moondream Python Client Library
+
+Python client library for moondream. This library is an alpha preview -- it is
+in an early stage of development, and backward compatibility is not yet
+guaranteed. If you are using this in production, please pin the revision you
+are using.
+
+This library currently offers optimized CPU inference, but will be slower than
+the PyTorch implementation for CUDA and MPS backends. If you are running on a
+Mac with M1/M2/M3 etc. chips, or if you have a GPU available, this library is
+not recommended yet.
+
+## Setup
+
+Install the library using pip:
+
+```
+pip install moondream==0.0.1
+```
+
+Then download the model weights:
+
+```
+# int8 weights (recommended):
+wget "https://huggingface.co/vikhyatk/moondream2/resolve/client/moondream-latest-int8.bin.gz?download=true" -O - | gunzip > moondream-latest-int8.bin
+# ...or, for FP16 weights:
+wget "https://huggingface.co/vikhyatk/moondream2/resolve/client/moondream-latest-f16.bin.gz?download=true" -O - | gunzip > moondream-latest-f16.bin
+```
+
+## Usage
+
+```python
+import moondream as md
+from PIL import Image
+
+model = md.VL("moondream-latest-int8.bin")
+image = Image.open("path/to/image.jpg")
+
+# Optional -- encode the image to efficiently run multiple queries on the same
+# image. This is not mandatory, since the model will automatically encode the
+# image if it is not already encoded.
+encoded_image = model.encode_image(image)
+
+# Caption the image.
+caption = model.caption(encoded_image)
+
+# ...or, if you want to stream the output:
+for t in model.caption(encoded_image, stream=True)["caption"]:
+    print(t, end="", flush=True)
+
+# Ask a question about the image.
+question = "How many people are in this image?"
+answer = model.answer_question(encoded_image, question)["answer"]
+
+# ...or again, if you want to stream the output:
+for t in model.answer_question(encoded_image, question, stream=True)["answer"]:
+    print(t, end="", flush=True)
+```
+
+### Accelerators
+
+(TK -- document how ONNX execution providers work.)
+

moondream-0.0.1.dist-info/RECORD

@@ -0,0 +1,6 @@
+moondream/__init__.py,sha256=gXJJSmxgNDWER5FA9OYQa8y87NtPxpJxnzCTXtFsMAE,19
+moondream/preprocess.py,sha256=19hE-Lzf_SbRLDJCS_LkX7Y7qoeHxivk_hF8fKbiufU,1970
+moondream/vl.py,sha256=ns1Rv18yzsOHsIkCvpSHvZv3dNqmKGB58DZHhHMiux0,9162
+moondream-0.0.1.dist-info/METADATA,sha256=8RJG1zGQInR1rd328z_WGBerPe1CEe7-oBDfEPyELR8,2656
+moondream-0.0.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+moondream-0.0.1.dist-info/RECORD,,

moondream-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.9.1
+Root-Is-Purelib: true
+Tag: py3-none-any