pixeltable 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

pixeltable/ext/functions/yolox.py

@@ -0,0 +1,92 @@
+import logging
+from pathlib import Path
+from typing import Iterable, Iterator
+from urllib.request import urlretrieve
+
+import PIL.Image
+import numpy as np
+import torch
+from yolox.data import ValTransform
+from yolox.exp import get_exp, Exp
+from yolox.models import YOLOX
+from yolox.utils import postprocess
+
+import pixeltable as pxt
+from pixeltable import env
+from pixeltable.func import Batch
+from pixeltable.functions.util import resolve_torch_device
+
+_logger = logging.getLogger('pixeltable')
+
+
+@pxt.udf(batch_size=4)
+def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
+    """
+    Runs the specified YOLOX object detection model on an image.
+
+    YOLOX support is part of the `pixeltable.ext` package: long-term support is not guaranteed, and it is not
+    intended for use in production applications.
+
+    Parameters:
+    - `model_id` - one of: `yolox_nano, `yolox_tiny`, `yolox_s`, `yolox_m`, `yolox_l`, `yolox_x`
+    - `threshold` - the threshold for object detection
+    """
+    model, exp = _lookup_model(model_id, 'cpu')
+    image_tensors = list(_images_to_tensors(images, exp))
+    batch_tensor = torch.stack(image_tensors)
+
+    with torch.no_grad():
+        output_tensor = model(batch_tensor)
+
+    outputs = postprocess(
+        output_tensor, 80, threshold, exp.nmsthre, class_agnostic=False
+    )
+
+    results: list[dict] = []
+    for image in images:
+        ratio = min(exp.test_size[0] / image.height, exp.test_size[1] / image.width)
+        if outputs[0] is None:
+            results.append({'bboxes': [], 'scores': [], 'labels': []})
+        else:
+            results.append({
+                'bboxes': [(output[:4] / ratio).tolist() for output in outputs[0]],
+                'scores': [output[4].item() * output[5].item() for output in outputs[0]],
+                'labels': [int(output[6]) for output in outputs[0]]
+            })
+    return results
+
+
+def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: Exp) -> Iterator[torch.Tensor]:
+    for image in images:
+        image_transform, _ = _val_transform(np.array(image), None, exp.test_size)
+        yield torch.from_numpy(image_transform)
+
+
+def _lookup_model(model_id: str, device: str) -> (YOLOX, Exp):
+    key = (model_id, device)
+    if key in _model_cache:
+        return _model_cache[key]
+
+    weights_url = f'https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/{model_id}.pth'
+    weights_file = Path(f'{env.Env.get().tmp_dir}/{model_id}.pth')
+    if not weights_file.exists():
+        _logger.info(f'Downloading weights for YOLOX model {model_id}: from {weights_url} -> {weights_file}')
+        urlretrieve(weights_url, weights_file)
+
+    exp = get_exp(exp_name=model_id)
+    model = exp.get_model().to(device)
+
+    model.eval()
+    model.head.training = False
+    model.training = False
+
+    # Load in the weights from training
+    weights = torch.load(weights_file, map_location=torch.device(device))
+    model.load_state_dict(weights['model'])
+
+    _model_cache[key] = (model, exp)
+    return model, exp
+
+
+_model_cache = {}
+_val_transform = ValTransform(legacy=False)

pixeltable/func/aggregate_function.py

@@ -3,13 +3,14 @@ from __future__ import annotations
 import abc
 import importlib
 import inspect
-from typing import Optional, Any, Type, List, Dict
+from typing import Optional, Any, Type, List, Dict, Callable
 import itertools
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 from .function import Function
 from .signature import Signature, Parameter
+from .globals import validate_symbol_path
 
 
 class Aggregator(abc.ABC):
@@ -136,8 +137,7 @@ def uda(
         update_types: List[ts.ColumnType],
         init_types: Optional[List[ts.ColumnType]] = None,
         requires_order_by: bool = False, allows_std_agg: bool = True, allows_window: bool = False,
-        name: Optional[str] = None
-) -> Type[Aggregator]:
+) -> Callable:
     """Decorator for user-defined aggregate functions.
 
     The decorated class must inherit from Aggregator and implement the following methods:
@@ -155,14 +155,11 @@ def uda(
     - requires_order_by: if True, the first parameter to the function is the order-by expression
     - allows_std_agg: if True, the function can be used as a standard aggregate function w/o a window
     - allows_window: if True, the function can be used with a window
-    - name: name of the AggregateFunction instance; if None, the class name is used
     """
-    if name is not None and not name.isidentifier():
-        raise excs.Error(f'Invalid name: {name}')
     if init_types is None:
         init_types = []
 
-    def decorator(cls: Type[Aggregator]) -> Type[Aggregator]:
+    def decorator(cls: Type[Aggregator]) -> Type[Function]:
         # validate type parameters
         num_init_params = len(inspect.signature(cls.__init__).parameters) - 1
         if num_init_params > 0:
@@ -178,17 +175,20 @@ def uda(
         assert value_type is not None
 
         # the AggregateFunction instance resides in the same module as cls
-        module_path = cls.__module__
-        nonlocal name
-        name = name or cls.__name__
-        instance_path = f'{module_path}.{name}'
+        class_path = f'{cls.__module__}.{cls.__qualname__}'
+        # nonlocal name
+        # name = name or cls.__name__
+        # instance_path_elements = class_path.split('.')[:-1] + [name]
+        # instance_path = '.'.join(instance_path_elements)
 
         # create the corresponding AggregateFunction instance
         instance = AggregateFunction(
-            cls, instance_path, init_types, update_types, value_type, requires_order_by, allows_std_agg, allows_window)
-        module = importlib.import_module(module_path)
-        setattr(module, name, instance)
+            cls, class_path, init_types, update_types, value_type, requires_order_by, allows_std_agg, allows_window)
+        # do the path validation at the very end, in order to be able to write tests for the other failure cases
+        validate_symbol_path(class_path)
+        #module = importlib.import_module(cls.__module__)
+        #setattr(module, name, instance)
 
-        return cls
+        return instance
 
     return decorator

pixeltable/func/expr_template_function.py

@@ -50,9 +50,17 @@ class ExprTemplateFunction(Function):
         bound_args.update(
             {param_name: default for param_name, default in self.defaults.items() if param_name not in bound_args})
         result = self.expr.copy()
+        import pixeltable.exprs as exprs
         for param_name, arg in bound_args.items():
             param_expr = self.param_exprs_by_name[param_name]
-            result = result.substitute(param_expr, arg)
+            if not isinstance(arg, exprs.Expr):
+                # TODO: use the available param_expr.col_type
+                arg_expr = exprs.Expr.from_object(arg)
+                if arg_expr is None:
+                    raise excs.Error(f'{self.self_name}(): cannot convert argument {arg} to a Pixeltable expression')
+            else:
+                arg_expr = arg
+            result = result.substitute(param_expr, arg_expr)
         import pixeltable.exprs as exprs
         assert not result.contains(exprs.Variable)
         return result

pixeltable/func/globals.py

@@ -1,29 +1,39 @@
-from typing import Optional
-from types import ModuleType
 import importlib
 import inspect
+from types import ModuleType
+from typing import Optional
 
+import pixeltable.exceptions as excs
 
-def resolve_symbol(symbol_path: str) -> object:
+
+def resolve_symbol(symbol_path: str) -> Optional[object]:
     path_elems = symbol_path.split('.')
     module: Optional[ModuleType] = None
-    if path_elems[0:2] == ['pixeltable', 'functions'] and len(path_elems) > 2:
-        # if this is a pixeltable.functions submodule, it cannot be resolved via pixeltable.functions;
-        # try to import the submodule directly
-        submodule_path = '.'.join(path_elems[0:3])
+    i = len(path_elems) - 1
+    while i > 0 and module is None:
         try:
-            module = importlib.import_module(submodule_path)
-            path_elems = path_elems[3:]
+            module = importlib.import_module('.'.join(path_elems[:i]))
         except ModuleNotFoundError:
-            pass
-    if module is None:
-        module = importlib.import_module(path_elems[0])
-        path_elems = path_elems[1:]
+            i -= 1
+    if i == 0:
+        return None  # Not resolvable
     obj = module
-    for el in path_elems:
+    for el in path_elems[i:]:
         obj = getattr(obj, el)
     return obj
 
+
+def validate_symbol_path(fn_path: str) -> None:
+    path_elems = fn_path.split('.')
+    fn_name = path_elems[-1]
+    if any(el == '<locals>' for el in path_elems):
+        raise excs.Error(
+            f'{fn_name}(): nested functions are not supported. Move the function to the module level or into a class.')
+    if any(not el.isidentifier() for el in path_elems):
+        raise excs.Error(
+            f'{fn_name}(): cannot resolve symbol path {fn_path}. Move the function to the module level or into a class.')
+
+
 def get_caller_module_path() -> str:
     """Return the module path of our caller's caller"""
     stack = inspect.stack()

pixeltable/func/signature.py

@@ -114,20 +114,12 @@ class Signature:
         return (col_type, is_batched)
 
     @classmethod
-    def create(
-            cls, c: Callable,
-            param_types: Optional[List[ts.ColumnType]] = None,
-            return_type: Optional[Union[ts.ColumnType, Callable]] = None
-    ) -> Signature:
-        """Create a signature for the given Callable.
-        Infer the parameter and return types, if none are specified.
-        Raises an exception if the types cannot be inferred.
-        """
+    def create_parameters(
+            cls, c: Callable, param_types: Optional[List[ts.ColumnType]] = None) -> List[Parameter]:
         sig = inspect.signature(c)
         py_parameters = list(sig.parameters.values())
-
-        # check non-var parameters for name collisions and default value compatibility
         parameters: List[Parameter] = []
+
         for idx, param in enumerate(py_parameters):
             if param.name in cls.SPECIAL_PARAM_NAMES:
                 raise excs.Error(f"'{param.name}' is a reserved parameter name")
@@ -135,6 +127,7 @@ class Signature:
                 parameters.append(Parameter(param.name, None, param.kind, False))
                 continue
 
+            # check non-var parameters for name collisions and default value compatibility
             if param_types is not None:
                 if idx >= len(param_types):
                     raise excs.Error(f'Missing type for parameter {param.name}')
@@ -155,7 +148,20 @@ class Signature:
 
             parameters.append(Parameter(param.name, param_type, param.kind, is_batched))
 
-        return_is_batched = False
+        return parameters
+
+    @classmethod
+    def create(
+            cls, c: Callable,
+            param_types: Optional[List[ts.ColumnType]] = None,
+            return_type: Optional[Union[ts.ColumnType, Callable]] = None
+    ) -> Signature:
+        """Create a signature for the given Callable.
+        Infer the parameter and return types, if none are specified.
+        Raises an exception if the types cannot be inferred.
+        """
+        parameters = cls.create_parameters(c, param_types)
+        sig = inspect.signature(c)
         if return_type is None:
             return_type, return_is_batched = cls._infer_type(sig.return_annotation)
             if return_type is None:

pixeltable/func/udf.py

@@ -11,6 +11,7 @@ from .callable_function import CallableFunction
 from .expr_template_function import ExprTemplateFunction
 from .function import Function
 from .function_registry import FunctionRegistry
+from .globals import validate_symbol_path
 from .signature import Signature
 
 
@@ -124,6 +125,8 @@ def make_function(
 
     # If this function is part of a module, register it
     if function_path is not None:
+        # do the validation at the very end, so it's easier to write tests for other failure scenarios
+        validate_symbol_path(function_path)
         FunctionRegistry.get().register_function(function_path, result)
 
     return result
@@ -142,17 +145,19 @@ def expr_udf(*args: Any, **kwargs: Any) -> Any:
         else:
             function_path = None
 
-        sig = Signature.create(py_fn, param_types=param_types, return_type=None)
         # TODO: verify that the inferred return type matches that of the template
         # TODO: verify that the signature doesn't contain batched parameters
 
         # construct Parameters from the function signature
+        params = Signature.create_parameters(py_fn, param_types=param_types)
         import pixeltable.exprs as exprs
-        var_exprs = [exprs.Variable(param.name, param.col_type) for param in sig.parameters.values()]
+        var_exprs = [exprs.Variable(param.name, param.col_type) for param in params]
         # call the function with the parameter expressions to construct an Expr with parameters
         template = py_fn(*var_exprs)
         assert isinstance(template, exprs.Expr)
         py_sig = inspect.signature(py_fn)
+        if function_path is not None:
+            validate_symbol_path(function_path)
         return ExprTemplateFunction(template, py_signature=py_sig, self_path=function_path, name=py_fn.__name__)
 
     if len(args) == 1:

pixeltable/functions/__init__.py

@@ -15,7 +15,7 @@ import pixeltable.functions.pil.image
 from pixeltable import exprs
 from pixeltable.type_system import IntType, ColumnType, FloatType, ImageType, VideoType
 # automatically import all submodules so that the udfs get registered
-from . import image, string, video, openai, together, fireworks, huggingface
+from . import image, string, video, huggingface
 
 # TODO: remove and replace calls with astype()
 def cast(expr: exprs.Expr, target_type: ColumnType) -> exprs.Expr:
@@ -23,8 +23,8 @@ def cast(expr: exprs.Expr, target_type: ColumnType) -> exprs.Expr:
     return expr
 
 @func.uda(
-    update_types=[IntType()], value_type=IntType(), name='sum', allows_window=True, requires_order_by=False)
-class SumAggregator(func.Aggregator):
+    update_types=[IntType()], value_type=IntType(), allows_window=True, requires_order_by=False)
+class sum(func.Aggregator):
     def __init__(self):
         self.sum: Union[int, float] = 0
     def update(self, val: Union[int, float]) -> None:
@@ -35,8 +35,8 @@ class SumAggregator(func.Aggregator):
 
 
 @func.uda(
-    update_types=[IntType()], value_type=IntType(), name='count', allows_window = True, requires_order_by = False)
-class CountAggregator(func.Aggregator):
+    update_types=[IntType()], value_type=IntType(), allows_window = True, requires_order_by = False)
+class count(func.Aggregator):
     def __init__(self):
         self.count = 0
     def update(self, val: int) -> None:
@@ -47,8 +47,8 @@ class CountAggregator(func.Aggregator):
 
 
 @func.uda(
-    update_types=[IntType()], value_type=FloatType(), name='mean', allows_window=False, requires_order_by=False)
-class MeanAggregator(func.Aggregator):
+    update_types=[IntType()], value_type=FloatType(), allows_window=False, requires_order_by=False)
+class mean(func.Aggregator):
     def __init__(self):
         self.sum = 0
         self.count = 0
@@ -63,9 +63,9 @@ class MeanAggregator(func.Aggregator):
 
 
 @func.uda(
-    init_types=[IntType()], update_types=[ImageType()], value_type=VideoType(), name='make_video',
+    init_types=[IntType()], update_types=[ImageType()], value_type=VideoType(),
     requires_order_by=True, allows_window=False)
-class VideoAggregator(func.Aggregator):
+class make_video(func.Aggregator):
     def __init__(self, fps: int = 25):
         """follows https://pyav.org/docs/develop/cookbook/numpy.html#generating-video"""
         self.container: Optional[av.container.OutputContainer] = None

pixeltable/functions/eval.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 from typing import List, Tuple, Dict
 from collections import defaultdict
 import sys
@@ -157,16 +156,16 @@ def calculate_image_tpfp(
         ts.JsonType(nullable=False)
     ])
 def eval_detections(
-        pred_bboxes: List[List[int]], pred_classes: List[int], pred_scores: List[float],
-        gt_bboxes: List[List[int]], gt_classes: List[int]
+        pred_bboxes: List[List[int]], pred_labels: List[int], pred_scores: List[float],
+        gt_bboxes: List[List[int]], gt_labels: List[int]
 ) -> Dict:
-    class_idxs = list(set(pred_classes + gt_classes))
+    class_idxs = list(set(pred_labels + gt_labels))
     result: List[Dict] = []
     pred_bboxes_arr = np.asarray(pred_bboxes)
-    pred_classes_arr = np.asarray(pred_classes)
+    pred_classes_arr = np.asarray(pred_labels)
     pred_scores_arr = np.asarray(pred_scores)
     gt_bboxes_arr = np.asarray(gt_bboxes)
-    gt_classes_arr = np.asarray(gt_classes)
+    gt_classes_arr = np.asarray(gt_labels)
     for class_idx in class_idxs:
         pred_filter = pred_classes_arr == class_idx
         gt_filter = gt_classes_arr == class_idx
@@ -181,8 +180,8 @@ def eval_detections(
     return result
 
 @func.uda(
-    update_types=[ts.JsonType()], value_type=ts.JsonType(), name='mean_ap', allows_std_agg=True, allows_window=False)
-class MeanAPAggregator:
+    update_types=[ts.JsonType()], value_type=ts.JsonType(), allows_std_agg=True, allows_window=False)
+class mean_ap(func.Aggregator):
     def __init__(self):
         self.class_tpfp: Dict[int, List[Dict]] = defaultdict(list)
 

pixeltable/functions/fireworks.py

@@ -1,61 +1,34 @@
-import logging
-import os
 from typing import Optional
 
+import fireworks.client
+
 import pixeltable as pxt
-import pixeltable.exceptions as excs
 from pixeltable import env
 
 
+def fireworks_client() -> fireworks.client.Fireworks:
+    return env.Env.get().get_client('fireworks', lambda api_key: fireworks.client.Fireworks(api_key=api_key))
+
+
 @pxt.udf
 def chat_completions(
-        prompt: str,
-        model: str,
+        messages: list[dict[str, str]],
         *,
+        model: str,
         max_tokens: Optional[int] = None,
-        repetition_penalty: Optional[float] = None,
         top_k: Optional[int] = None,
         top_p: Optional[float] = None,
         temperature: Optional[float] = None
 ) -> dict:
-    initialize()
     kwargs = {
         'max_tokens': max_tokens,
-        'repetition_penalty': repetition_penalty,
         'top_k': top_k,
         'top_p': top_p,
         'temperature': temperature
     }
     kwargs_not_none = dict(filter(lambda x: x[1] is not None, kwargs.items()))
-    import fireworks.client
-    return fireworks.client.Completion.create(
+    return fireworks_client().chat.completions.create(
         model=model,
-        prompt_or_messages=prompt,
+        messages=messages,
         **kwargs_not_none
     ).dict()
-
-
-def initialize():
-    global _is_fireworks_initialized
-    if _is_fireworks_initialized:
-        return
-
-    _logger.info('Initializing Fireworks client.')
-
-    config = pxt.env.Env.get().config
-
-    if 'fireworks' in config and 'api_key' in config['fireworks']:
-        api_key = config['fireworks']['api_key']
-    else:
-        api_key = os.environ.get('FIREWORKS_API_KEY')
-    if api_key is None or api_key == '':
-        raise excs.Error('Fireworks client not initialized (no API key configured).')
-
-    import fireworks.client
-
-    fireworks.client.api_key = api_key
-    _is_fireworks_initialized = True
-
-
-_logger = logging.getLogger('pixeltable')
-_is_fireworks_initialized = False

pixeltable/functions/huggingface.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable
+from typing import Callable, TypeVar, Optional
 
 import PIL.Image
 import numpy as np
@@ -7,10 +7,13 @@ import pixeltable as pxt
 import pixeltable.env as env
 import pixeltable.type_system as ts
 from pixeltable.func import Batch
+from pixeltable.functions.util import resolve_torch_device
 
 
 @pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
-def sentence_transformer(sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False) -> Batch[np.ndarray]:
+def sentence_transformer(
+        sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False
+) -> Batch[np.ndarray]:
     env.Env.get().require_package('sentence_transformers')
     from sentence_transformers import SentenceTransformer
 
@@ -53,44 +56,60 @@ def cross_encoder_list(sentence1: str, sentences2: list, *, model_id: str) -> li
     return array.tolist()
 
 
-@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((512,), dtype=ts.FloatType(), nullable=False))
 def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
     env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
     from transformers import CLIPModel, CLIPProcessor
 
-    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
+    assert model.config.projection_dim == 512
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
-    inputs = processor(text=text, return_tensors='pt', padding=True, truncation=True)
-    embeddings = model.get_text_features(**inputs).detach().numpy()
+    with torch.no_grad():
+        inputs = processor(text=text, return_tensors='pt', padding=True, truncation=True)
+        embeddings = model.get_text_features(**inputs.to(device)).detach().to('cpu').numpy()
+
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
-@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((512,), dtype=ts.FloatType(), nullable=False))
 def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndarray]:
     env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
     from transformers import CLIPModel, CLIPProcessor
 
-    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
+    assert model.config.projection_dim == 512
     processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
 
-    inputs = processor(images=image, return_tensors='pt', padding=True)
-    embeddings = model.get_image_features(**inputs).detach().numpy()
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors='pt', padding=True)
+        embeddings = model.get_image_features(**inputs.to(device)).detach().to('cpu').numpy()
+
     return [embeddings[i] for i in range(embeddings.shape[0])]
 
 
-@pxt.udf(batch_size=32)
+@pxt.udf(batch_size=4)
 def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
     env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
     from transformers import DetrImageProcessor, DetrForObjectDetection
 
-    model = _lookup_model(model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'))
+    model = _lookup_model(
+        model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'), device=device)
     processor = _lookup_processor(model_id, lambda x: DetrImageProcessor.from_pretrained(x, revision='no_timm'))
 
-    inputs = processor(images=image, return_tensors='pt')
-    outputs = model(**inputs)
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors='pt')
+        outputs = model(**inputs.to(device))
+        results = processor.post_process_object_detection(
+            outputs, threshold=threshold, target_sizes=[(img.height, img.width) for img in image]
+        )
 
-    results = processor.post_process_object_detection(outputs, threshold=threshold)
     return [
         {
             'scores': [score.item() for score in result['scores']],
@@ -102,14 +121,23 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
     ]
 
 
-def _lookup_model(model_id: str, create: Callable) -> Any:
-    key = (model_id, create)  # For safety, include the `create` callable in the cache key
+T = TypeVar('T')
+
+
+def _lookup_model(model_id: str, create: Callable[[str], T], device: Optional[str] = None) -> T:
+    from torch import nn
+    key = (model_id, create, device)  # For safety, include the `create` callable in the cache key
     if key not in _model_cache:
-        _model_cache[key] = create(model_id)
+        model = create(model_id)
+        if device is not None:
+            model.to(device)
+        if isinstance(model, nn.Module):
+            model.eval()
+        _model_cache[key] = model
     return _model_cache[key]
 
 
-def _lookup_processor(model_id: str, create: Callable) -> Any:
+def _lookup_processor(model_id: str, create: Callable[[str], T]) -> T:
     key = (model_id, create)  # For safety, include the `create` callable in the cache key
     if key not in _processor_cache:
         _processor_cache[key] = create(model_id)