corvic-engine 0.3.0rc62__cp38-abi3-win_amd64.whl → 0.3.0rc63__cp38-abi3-win_amd64.whl

corvic/system/_embedder.py

@@ -1,6 +1,6 @@
 import dataclasses
 from collections.abc import Sequence
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import polars as pl
@@ -9,6 +9,12 @@ from typing_extensions import Protocol
 from corvic import orm
 from corvic.result import InternalError, InvalidArgumentError, Ok
 
+if TYPE_CHECKING:
+    from transformers import (
+        CLIPModel,
+        CLIPProcessor,
+    )
+
 
 @dataclasses.dataclass
 class EmbedTextContext:
@@ -64,6 +70,12 @@ class ImageEmbedder(Protocol):
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError: ...
 
 
+@dataclasses.dataclass
+class ClipModels:
+    model: "CLIPModel"
+    processor: "CLIPProcessor"
+
+
 class ClipText(TextEmbedder):
     """Clip Text embedder.
 
@@ -76,28 +88,39 @@ class ClipText(TextEmbedder):
     overcoming several major challenges in computer vision.
     """
 
-    def embed(
-        self, context: EmbedTextContext
-    ) -> Ok[EmbedTextResult] | InvalidArgumentError | InternalError:
-        import torch
+    def _load_models(self):
         from transformers import (
             CLIPModel,
             CLIPProcessor,
         )
 
         model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-            "openai/clip-vit-base-patch32"
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
         )
         processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
-            "openai/clip-vit-base-patch32"
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+            use_fast=False,
         )
-        model.eval()
+        return ClipModels(model=model, processor=processor)
+
+    def embed(
+        self, context: EmbedTextContext
+    ) -> Ok[EmbedTextResult] | InvalidArgumentError | InternalError:
         match context.expected_coordinate_bitwidth:
             case 64:
                 coord_dtype = pl.Float64()
             case 32:
                 coord_dtype = pl.Float32()
 
+        models = self._load_models()
+        model = models.model
+        processor = models.processor
+        model.eval()
+
+        import torch
+
         with torch.no_grad():
             inputs: dict[str, torch.Tensor] = processor(  # pyright: ignore[reportAssignmentType]
                 text=context.inputs,

corvic/system/_image_embedder.py

@@ -1,3 +1,4 @@
+import dataclasses
 from io import BytesIO
 from typing import TYPE_CHECKING, Any
 
@@ -13,6 +14,10 @@ from corvic.system._embedder import (
 
 if TYPE_CHECKING:
     from PIL import Image
+    from transformers import (
+        CLIPModel,
+        CLIPProcessor,
+    )
 
 
 class RandomImageEmbedder(ImageEmbedder):
@@ -58,6 +63,12 @@ def image_from_bytes(
         return InvalidArgumentError("invalid image format")
 
 
+@dataclasses.dataclass
+class ClipModels:
+    model: "CLIPModel"
+    processor: "CLIPProcessor"
+
+
 class Clip(ImageEmbedder):
     """Clip image embedder.
 
@@ -70,6 +81,23 @@ class Clip(ImageEmbedder):
     overcoming several major challenges in computer vision.
     """
 
+    def _load_models(self):
+        from transformers import (
+            CLIPModel,
+            CLIPProcessor,
+        )
+
+        model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+        )
+        processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+            use_fast=False,
+        )
+        return ClipModels(model=model, processor=processor)
+
     def embed(
         self, context: EmbedImageContext
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError:
@@ -99,20 +127,13 @@ class Clip(ImageEmbedder):
                 )
             )
 
-        import torch
-        from transformers import (
-            CLIPModel,
-            CLIPProcessor,
-        )
-
-        model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-            "openai/clip-vit-base-patch32"
-        )
-        processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
-            "openai/clip-vit-base-patch32"
-        )
+        models = self._load_models()
+        model = models.model
+        processor = models.processor
         model.eval()
 
+        import torch
+
         with torch.no_grad():
             inputs: dict[str, torch.FloatTensor] = processor(  # pyright: ignore[reportAssignmentType]
                 images=images, return_tensors="pt"