inference-models 0.18.3__py3-none-any.whl

inference_models/models/sam2_rt/sam2_pytorch.py

@@ -0,0 +1,119 @@
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.errors import MissingDependencyError, ModelRuntimeError
+from inference_models.models.common.model_packages import get_model_package_contents
+
+try:
+    import hydra
+    from sam2.build_sam import build_sam2_camera_predictor
+    from sam2.sam2_camera_predictor import SAM2CameraPredictor
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message=f"Could not import SAM2 model, please consult README for installation instructions.",
+    ) from import_error
+
+
+class SAM2ForStream:
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        **kwargs,
+    ) -> "SAM2ForStream":
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=[
+                "weights.pt",
+                "sam2-rt.yaml",
+            ],
+        )
+        hydra.core.global_hydra.GlobalHydra.instance().clear()
+        hydra.initialize_config_dir(
+            config_dir=Path(model_package_content["sam2-rt.yaml"]).parent.as_posix(),
+            version_base=None,
+        )
+        predictor: SAM2CameraPredictor = build_sam2_camera_predictor(
+            config_file=Path(model_package_content["sam2-rt.yaml"]).name,
+            ckpt_path=model_package_content["weights.pt"],
+            device=device,
+        )
+        return cls(predictor=predictor, device=device)
+
+    def __init__(self, predictor: SAM2CameraPredictor, device: torch.device):
+        self._predictor = predictor
+        self._device = device
+
+    def prompt(
+        self,
+        image: Union[np.ndarray, torch.Tensor],
+        bboxes: Union[Tuple[int, int, int, int], List[Tuple[int, int, int, int]]],
+        state_dict: Optional[dict] = None,
+        clear_old_points: bool = True,
+        normalize_coords: bool = True,
+        frame_idx: int = 0,
+    ) -> tuple:
+        if isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        if clear_old_points or not self._predictor.condition_state:
+            self._predictor.load_first_frame(image)
+        if state_dict is not None:
+            self._predictor.load_state_dict(state_dict)
+        obj_id = 0
+        if (
+            self._predictor.condition_state
+            and self._predictor.condition_state["obj_ids"]
+        ):
+            obj_id = max(self._predictor.condition_state["obj_ids"]) + 1
+        if not isinstance(bboxes, list):
+            bboxes = [bboxes]
+        for pts in bboxes:
+            if len(pts) < 4:
+                continue
+            x1, y1, x2, y2 = pts[:4]
+            x_lt = int(round(min(x1, x2)))
+            y_lt = int(round(min(y1, y2)))
+            x_rb = int(round(max(x1, x2)))
+            y_rb = int(round(max(y1, y2)))
+            xyxy = np.array([[x_lt, y_lt, x_rb, y_rb]])
+
+            _, object_ids, mask_logits = self._predictor.add_new_prompt(
+                frame_idx=frame_idx,
+                obj_id=obj_id,
+                bbox=xyxy,
+                clear_old_points=clear_old_points,
+                normalize_coords=normalize_coords,
+            )
+            obj_id += 1
+        masks = (mask_logits > 0.0).cpu().numpy()
+        masks = np.squeeze(masks).astype(bool)
+        if len(masks.shape) == 2:
+            masks = np.expand_dims(masks, axis=0)
+        object_ids = np.array(object_ids)
+        return masks, object_ids, self._predictor.state_dict()
+
+    def track(
+        self,
+        image: Union[np.ndarray, torch.Tensor],
+        state_dict: Optional[dict] = None,
+    ) -> tuple:
+        if isinstance(image, torch.Tensor):
+            image = image.detach().cpu().numpy()
+        if state_dict is not None:
+            self._predictor.load_state_dict(state_dict)
+        if not self._predictor.condition_state:
+            raise ModelRuntimeError(
+                "Attempt to track with no prior call to prompt; prompt must be called first"
+            )
+        object_ids, mask_logits = self._predictor.track(image)
+        masks = (mask_logits > 0.0).cpu().numpy()
+        masks = np.squeeze(masks).astype(bool)
+        if len(masks.shape) == 2:
+            masks = np.expand_dims(masks, axis=0)
+        object_ids = np.array(object_ids)
+        return masks, object_ids, self._predictor.state_dict()

inference_models/models/smolvlm/smolvlm_hf.py

@@ -0,0 +1,245 @@
+import os
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from peft import PeftModel
+from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig
+
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat
+from inference_models.models.common.roboflow.model_packages import (
+    InferenceConfig,
+    ResizeMode,
+    parse_inference_config,
+)
+from inference_models.models.common.roboflow.pre_processing import (
+    pre_process_network_input,
+)
+
+
+class SmolVLMHF:
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        trust_remote_code: bool = False,
+        local_files_only: bool = True,
+        quantization_config: Optional[BitsAndBytesConfig] = None,
+        disable_quantization: bool = False,
+        **kwargs,
+    ) -> "SmolVLMHF":
+        torch_dtype = torch.float16 if device.type == "cuda" else torch.float32
+        inference_config_path = os.path.join(
+            model_name_or_path, "inference_config.json"
+        )
+        inference_config = None
+        if os.path.exists(inference_config_path):
+            inference_config = parse_inference_config(
+                config_path=inference_config_path,
+                allowed_resize_modes={
+                    ResizeMode.STRETCH_TO,
+                    ResizeMode.LETTERBOX,
+                    ResizeMode.CENTER_CROP,
+                    ResizeMode.LETTERBOX_REFLECT_EDGES,
+                },
+            )
+        adapter_config_path = os.path.join(model_name_or_path, "adapter_config.json")
+        if (
+            quantization_config is None
+            and device.type == "cuda"
+            and not disable_quantization
+        ):
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_quant_type="nf4",
+            )
+        if os.path.exists(adapter_config_path):
+
+            base_model_path = os.path.join(model_name_or_path, "base")
+            model = AutoModelForImageTextToText.from_pretrained(
+                base_model_path,
+                dtype=torch_dtype,
+                trust_remote_code=trust_remote_code,
+                local_files_only=local_files_only,
+                quantization_config=quantization_config,
+            )
+            model = PeftModel.from_pretrained(model, model_name_or_path)
+            if quantization_config is None:
+                model.merge_and_unload()
+            model.to(device)
+
+            processor = AutoProcessor.from_pretrained(
+                base_model_path,
+                padding_side="left",
+                trust_remote_code=trust_remote_code,
+                local_files_only=local_files_only,
+                use_fast=True,
+            )
+        else:
+            model = AutoModelForImageTextToText.from_pretrained(
+                model_name_or_path,
+                dtype=torch_dtype,
+                device_map=device,
+                trust_remote_code=trust_remote_code,
+                local_files_only=local_files_only,
+                quantization_config=quantization_config,
+            ).eval()
+            processor = AutoProcessor.from_pretrained(
+                model_name_or_path,
+                padding_side="left",
+                trust_remote_code=trust_remote_code,
+                local_files_only=local_files_only,
+                use_fast=True,
+            )
+        return cls(
+            model=model,
+            processor=processor,
+            inference_config=inference_config,
+            device=device,
+            torch_dtype=torch_dtype,
+        )
+
+    def __init__(
+        self,
+        model: AutoModelForImageTextToText,
+        processor: AutoProcessor,
+        inference_config: Optional[InferenceConfig],
+        device: torch.device,
+        torch_dtype: torch.dtype,
+    ):
+        self._model = model
+        self._processor = processor
+        self._inference_config = inference_config
+        self._device = device
+        self._torch_dtype = torch_dtype
+
+    def prompt(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        prompt: str,
+        images_to_single_prompt: bool = True,
+        input_color_format: Optional[ColorFormat] = None,
+        max_new_tokens: int = 400,
+        do_sample: bool = False,
+        skip_special_tokens: bool = True,
+        **kwargs,
+    ) -> List[str]:
+        inputs = self.pre_process_generation(
+            images=images,
+            prompt=prompt,
+            images_to_single_prompt=images_to_single_prompt,
+            input_color_format=input_color_format,
+        )
+        generated_ids = self.generate(
+            inputs=inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+        )
+        return self.post_process_generation(
+            generated_ids=generated_ids,
+            skip_special_tokens=skip_special_tokens,
+        )
+
+    def pre_process_generation(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        prompt: str,
+        images_to_single_prompt: bool = True,
+        input_color_format: Optional[ColorFormat] = None,
+        image_size: Optional[Tuple[int, int]] = None,
+        **kwargs,
+    ) -> dict:
+        def _to_tensor(image: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+            is_numpy = isinstance(image, np.ndarray)
+            if is_numpy:
+                tensor_image = torch.from_numpy(image.copy()).permute(2, 0, 1)
+            else:
+                tensor_image = image
+            if input_color_format == "bgr" or (is_numpy and input_color_format is None):
+                tensor_image = tensor_image[[2, 1, 0], :, :]
+            if image_size is not None:
+                tensor_image = torch.nn.functional.interpolate(
+                    image,
+                    [image_size[1], image_size[0]],
+                    mode="bilinear",
+                )
+            return tensor_image
+
+        if self._inference_config is None:
+            if isinstance(images, torch.Tensor) and images.ndim > 3:
+                image_list = [_to_tensor(img) for img in images]
+            elif not isinstance(images, list):
+                image_list = [_to_tensor(images)]
+            else:
+                image_list = [_to_tensor(img) for img in images]
+        else:
+            images = pre_process_network_input(
+                images=images,
+                image_pre_processing=self._inference_config.image_pre_processing,
+                network_input=self._inference_config.network_input,
+                target_device=self._device,
+                input_color_format=input_color_format,
+                image_size_wh=image_size,
+            )[0]
+            image_list = [e[0] for e in torch.split(images, 1, dim=0)]
+        if images_to_single_prompt:
+            content = [{"type": "image"}] * len(image_list)
+            content.append({"type": "text", "text": prompt})
+            conversations = [[{"role": "user", "content": content}]]
+        else:
+            conversations = []
+            for _ in image_list:
+                conversations.append(
+                    [
+                        {
+                            "role": "user",
+                            "content": [
+                                {"type": "image"},
+                                {"type": "text", "text": prompt},
+                            ],
+                        }
+                    ]
+                )
+        text_prompts = self._processor.apply_chat_template(
+            conversations, add_generation_prompt=True
+        )
+        max_image_size = None
+        if image_size:
+            max_image_size = {"longest_edge": max(image_size[0], image_size[1])}
+
+        inputs = self._processor(
+            text=text_prompts,
+            images=image_list,
+            return_tensors="pt",
+            padding=True,
+            max_image_size=max_image_size,
+        )
+        return inputs.to(self._device, dtype=self._torch_dtype)
+
+    def generate(
+        self,
+        inputs: dict,
+        max_new_tokens: int = 400,
+        do_sample: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        generation = self._model.generate(
+            **inputs, do_sample=do_sample, max_new_tokens=max_new_tokens
+        )
+        input_len = inputs["input_ids"].shape[-1]
+        return generation[:, input_len:]
+
+    def post_process_generation(
+        self,
+        generated_ids: torch.Tensor,
+        skip_special_tokens: bool = False,
+        **kwargs,
+    ) -> List[str]:
+        decoded = self._processor.batch_decode(
+            generated_ids, skip_special_tokens=skip_special_tokens
+        )
+        return [result.strip() for result in decoded]

inference_models/models/trocr/trocr_hf.py

@@ -0,0 +1,53 @@
+from typing import List, Union
+
+import numpy as np
+import torch
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.models.base.documents_parsing import TextOnlyOCRModel
+
+
+class TROcrHF(TextOnlyOCRModel[torch.Tensor, torch.Tensor]):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        local_files_only: bool = True,
+        **kwargs,
+    ) -> "TextOnlyOCRModel":
+        model = VisionEncoderDecoderModel.from_pretrained(
+            model_name_or_path,
+            local_files_only=local_files_only,
+        ).to(device)
+        processor = TrOCRProcessor.from_pretrained(
+            model_name_or_path, local_files_only=local_files_only
+        )
+        return cls(model=model, processor=processor, device=device)
+
+    def __init__(
+        self,
+        processor: TrOCRProcessor,
+        model: VisionEncoderDecoderModel,
+        device: torch.device,
+    ):
+        self._processor = processor
+        self._model = model
+        self._device = device
+
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        **kwargs,
+    ) -> torch.Tensor:
+        inputs = self._processor(images=images, return_tensors="pt")
+        return inputs["pixel_values"].to(self._device)
+
+    def forward(self, pre_processed_images: torch.Tensor, **kwargs) -> torch.Tensor:
+        with torch.inference_mode():
+            return self._model.generate(pre_processed_images)
+
+    def post_process(self, model_results: torch.Tensor, **kwargs) -> List[str]:
+        return self._processor.batch_decode(model_results, skip_special_tokens=True)