fastembed-gpu 0.7.4__tar.gz → 0.8.0__tar.gz

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fastembed-gpu
-Version: 0.7.4
+Version: 0.8.0
 Summary: Fast, light, accurate library built for retrieval embedding generation
 License: Apache License
 License-File: LICENSE
@@ -8,10 +8,9 @@ License-File: NOTICE
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
 Author: Qdrant Team
 Author-email: info@qdrant.tech
-Requires-Python: >=3.9.0
+Requires-Python: >=3.10.0
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -21,17 +20,17 @@ Requires-Dist: huggingface-hub (>=0.20,<2.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
 Requires-Dist: mmh3 (>=4.1.0,<6.0.0)
 Requires-Dist: numpy (>=1.21) ; python_version == "3.11"
-Requires-Dist: numpy (>=1.21,<2.1.0) ; python_version < "3.10"
 Requires-Dist: numpy (>=1.21,<2.3.0) ; python_version == "3.10"
 Requires-Dist: numpy (>=1.26) ; python_version == "3.12"
 Requires-Dist: numpy (>=2.1.0) ; python_version == "3.13"
 Requires-Dist: numpy (>=2.3.0) ; python_version >= "3.14"
-Requires-Dist: onnxruntime-gpu (>1.20.0) ; python_version >= "3.13"
-Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0) ; python_version >= "3.10" and python_version < "3.13"
-Requires-Dist: onnxruntime-gpu (>=1.17.0,<1.20.0) ; python_version < "3.10"
-Requires-Dist: pillow (>=10.3.0,<11.0) ; python_version < "3.10"
-Requires-Dist: pillow (>=10.3.0,<12.0) ; python_version >= "3.10" and python_version < "3.13"
-Requires-Dist: pillow (>=11.0.0,<12.0) ; python_version >= "3.13"
+Requires-Dist: onnxruntime-gpu (>1.21.0,!=1.24.0,!=1.24.1) ; python_version == "3.13"
+Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0,!=1.24.0,!=1.24.1) ; python_version >= "3.11" and python_version < "3.13"
+Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0,<1.24) ; python_version == "3.10"
+Requires-Dist: onnxruntime-gpu (>=1.24.2) ; python_version >= "3.14"
+Requires-Dist: pillow (>=10.3.0,<13.0) ; python_version >= "3.10" and python_version < "3.13"
+Requires-Dist: pillow (>=11.0.0,<13.0) ; python_version == "3.13"
+Requires-Dist: pillow (>=12.0.0,<13.0) ; python_version >= "3.14"
 Requires-Dist: py-rust-stemmers (>=0.1.0,<0.2.0)
 Requires-Dist: requests (>=2.31,<3.0)
 Requires-Dist: tokenizers (>=0.15,<1.0)

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/common/model_description.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Optional, Any
+from typing import Any
 
 
 @dataclass(frozen=True)
 class ModelSource:
-    hf: Optional[str] = None
-    url: Optional[str] = None
+    hf: str | None = None
+    url: str | None = None
     _deprecated_tar_struct: bool = False
 
     @property
@@ -33,8 +33,8 @@ class BaseModelDescription:
 
 @dataclass(frozen=True)
 class DenseModelDescription(BaseModelDescription):
-    dim: Optional[int] = None
-    tasks: Optional[dict[str, Any]] = field(default_factory=dict)
+    dim: int | None = None
+    tasks: dict[str, Any] | None = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         assert self.dim is not None, "dim is required for dense model description"
@@ -42,8 +42,8 @@ class DenseModelDescription(BaseModelDescription):
 
 @dataclass(frozen=True)
 class SparseModelDescription(BaseModelDescription):
-    requires_idf: Optional[bool] = None
-    vocab_size: Optional[int] = None
+    requires_idf: bool | None = None
+    vocab_size: int | None = None
 
 
 class PoolingType(str, Enum):

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/common/model_management.py

@@ -5,7 +5,7 @@ import shutil
 import tarfile
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Optional, Union, TypeVar, Generic
+from typing import Any, TypeVar, Generic
 
 import requests
 from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -180,8 +180,8 @@ class ModelManagement(Generic[T]):
 
         def _collect_file_metadata(
             model_dir: Path, repo_files: list[RepoFile]
-        ) -> dict[str, dict[str, Union[int, str]]]:
-            meta: dict[str, dict[str, Union[int, str]]] = {}
+        ) -> dict[str, dict[str, int | str]]:
+            meta: dict[str, dict[str, int | str]] = {}
             file_info_map = {f.path: f for f in repo_files}
             for file_path in model_dir.rglob("*"):
                 if file_path.is_file() and file_path.name != cls.METADATA_FILE:
@@ -193,9 +193,7 @@ class ModelManagement(Generic[T]):
                         }
             return meta
 
-        def _save_file_metadata(
-            model_dir: Path, meta: dict[str, dict[str, Union[int, str]]]
-        ) -> None:
+        def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int | str]]) -> None:
             try:
                 if not model_dir.exists():
                     model_dir.mkdir(parents=True, exist_ok=True)
@@ -397,7 +395,11 @@ class ModelManagement(Generic[T]):
             Path: The path to the downloaded model directory.
         """
         local_files_only = kwargs.get("local_files_only", False)
-        specific_model_path: Optional[str] = kwargs.pop("specific_model_path", None)
+        hf_offline = os.environ.get("HF_HUB_OFFLINE", "").strip().upper()
+        if not local_files_only and hf_offline in {"1", "TRUE", "YES", "ON"}:
+            local_files_only = True
+            kwargs["local_files_only"] = True
+        specific_model_path: str | None = kwargs.pop("specific_model_path", None)
         if specific_model_path:
             return Path(specific_model_path)
         retries = 1 if local_files_only else retries

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/common/onnx_model.py

@@ -1,7 +1,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
+from typing import Any, Generic, Iterable, Sequence, Type, TypeVar
 
 import numpy as np
 import onnxruntime as ort
@@ -9,7 +9,7 @@ import onnxruntime as ort
 from numpy.typing import NDArray
 from tokenizers import Tokenizer
 
-from fastembed.common.types import OnnxProvider, NumpyArray
+from fastembed.common.types import OnnxProvider, NumpyArray, Device
 from fastembed.parallel_processor import Worker
 
 # Holds type of the embedding result
@@ -19,8 +19,9 @@ T = TypeVar("T")
 @dataclass
 class OnnxOutputContext:
     model_output: NumpyArray
-    attention_mask: Optional[NDArray[np.int64]] = None
-    input_ids: Optional[NDArray[np.int64]] = None
+    attention_mask: NDArray[np.int64] | None = None
+    input_ids: NDArray[np.int64] | None = None
+    metadata: dict[str, Any] | None = None
 
 
 class OnnxModel(Generic[T]):
@@ -43,8 +44,8 @@ class OnnxModel(Generic[T]):
         raise NotImplementedError("Subclasses must implement this method")
 
     def __init__(self) -> None:
-        self.model: Optional[ort.InferenceSession] = None
-        self.tokenizer: Optional[Tokenizer] = None
+        self.model: ort.InferenceSession | None = None
+        self.tokenizer: Tokenizer | None = None
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -58,25 +59,30 @@ class OnnxModel(Generic[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
+        available_providers = ort.get_available_providers()
+        cuda_available = "CUDAExecutionProvider" in available_providers
+        explicit_cuda = cuda is True or cuda == Device.CUDA
 
-        if cuda and providers is not None:
+        if explicit_cuda and providers is not None:
             warnings.warn(
-                f"`cuda` and `providers` are mutually exclusive parameters, cuda: {cuda}, providers: {providers}",
+                f"`cuda` and `providers` are mutually exclusive parameters, "
+                f"cuda: {cuda}, providers: {providers}. If you'd like to use providers, cuda should be one of "
+                f"[False, Device.CPU, Device.AUTO].",
                 category=UserWarning,
                 stacklevel=6,
             )
 
         if providers is not None:
             onnx_providers = list(providers)
-        elif cuda:
+        elif explicit_cuda or (cuda == Device.AUTO and cuda_available):
             if device_id is None:
                 onnx_providers = ["CUDAExecutionProvider"]
             else:
@@ -84,7 +90,6 @@ class OnnxModel(Generic[T]):
         else:
             onnx_providers = ["CPUExecutionProvider"]
 
-        available_providers = ort.get_available_providers()
         requested_provider_names: list[str] = []
         for provider in onnx_providers:
             # check providers available

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/common/preprocessor_utils.py

@@ -50,9 +50,10 @@ def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=max_context)
-    tokenizer.enable_padding(
-        pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
-    )
+    if not tokenizer.padding:
+        tokenizer.enable_padding(
+            pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
+        )
 
     for token in tokens_map.values():
         if isinstance(token, str):

fastembed_gpu-0.8.0/fastembed/common/types.py

@@ -0,0 +1,27 @@
+from enum import Enum
+from pathlib import Path
+from typing import Any, TypeAlias
+
+import numpy as np
+from numpy.typing import NDArray
+from PIL import Image
+
+
+class Device(str, Enum):
+    CPU = "cpu"
+    CUDA = "cuda"
+    AUTO = "auto"
+
+
+PathInput: TypeAlias = str | Path
+ImageInput: TypeAlias = PathInput | Image.Image
+
+OnnxProvider: TypeAlias = str | tuple[str, dict[Any, Any]]
+NumpyArray: TypeAlias = (
+    NDArray[np.float64]
+    | NDArray[np.float32]
+    | NDArray[np.float16]
+    | NDArray[np.int8]
+    | NDArray[np.int64]
+    | NDArray[np.int32]
+)

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/common/utils.py

@@ -5,7 +5,7 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, Optional, TypeVar
+from typing import Iterable, TypeVar
 
 import numpy as np
 from numpy.typing import NDArray
@@ -45,7 +45,7 @@ def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
         yield b
 
 
-def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
+def define_cache_dir(cache_dir: str | None = None) -> Path:
     """
     Define the cache directory for fastembed
     """

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/embedding.py

@@ -1,4 +1,4 @@
-from typing import Optional, Any
+from typing import Any
 
 from loguru import logger
 
@@ -17,8 +17,8 @@ class JinaEmbedding(TextEmbedding):
     def __init__(
         self,
         model_name: str = "jinaai/jina-embeddings-v2-base-en",
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/image/image_embedding.py

@@ -1,7 +1,7 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 from dataclasses import asdict
 
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_embedding import OnnxImageEmbedding
@@ -48,11 +48,11 @@ class ImageEmbedding(ImageEmbeddingBase):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
         **kwargs: Any,
     ):
@@ -98,7 +98,7 @@ class ImageEmbedding(ImageEmbeddingBase):
             ValueError: If the model name is not found in the supported models.
         """
         descriptions = cls._list_supported_models()
-        embedding_size: Optional[int] = None
+        embedding_size: int | None = None
         for description in descriptions:
             if description.model.lower() == model_name.lower():
                 embedding_size = description.dim
@@ -113,9 +113,9 @@ class ImageEmbedding(ImageEmbeddingBase):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/image/image_embedding_base.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Optional, Any, Union
+from typing import Iterable, Any
 
 from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
@@ -10,21 +10,21 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
         **kwargs: Any,
     ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
-        self._embedding_size: Optional[int] = None
+        self._embedding_size: int | None = None
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/image/onnx_embedding.py

@@ -1,7 +1,7 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
@@ -63,14 +63,15 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
     def __init__(
         self,
         model_name: str,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         lazy_load: bool = False,
-        device_id: Optional[int] = None,
-        specific_model_path: Optional[str] = None,
+        device_id: int | None = None,
+        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         """
@@ -82,10 +83,11 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
             providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
-            cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
-                Defaults to False.
+            cuda (Union[bool, Device], optional): Whether to use cuda for inference. Mutually exclusive with `providers`
+                Defaults to Device.AUTO.
             device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
-                workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
+                workers. Should be used with `cuda` equals to `True`, `Device.AUTO` or `Device.CUDA`, mutually exclusive
+                with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
             device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
@@ -105,7 +107,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         self.cuda = cuda
 
         # This device_id will be used if we need to load model in current process
-        self.device_id: Optional[int] = None
+        self.device_id: int | None = None
         if device_id is not None:
             self.device_id = device_id
         elif self.device_ids is not None:
@@ -150,9 +152,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
     def embed(
         self,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 16,
-        parallel: Optional[int] = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
         """

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/image/onnx_image_model.py

@@ -2,13 +2,13 @@ import contextlib
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from PIL import Image
 
 from fastembed.image.transform.operators import Compose
-from fastembed.common.types import NumpyArray
+from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
@@ -37,7 +37,7 @@ class OnnxImageModel(OnnxModel[T]):
 
     def __init__(self) -> None:
         super().__init__()
-        self.processor: Optional[Compose] = None
+        self.processor: Compose | None = None
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -51,11 +51,11 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_dir: Path,
         model_file: str,
-        threads: Optional[int],
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_id: Optional[int] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
     ) -> None:
         super()._load_onnx_model(
             model_dir=model_dir,
@@ -76,9 +76,11 @@ class OnnxImageModel(OnnxModel[T]):
         return {input_name: encoded}
 
     def onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
-        with contextlib.ExitStack():
+        with contextlib.ExitStack() as stack:
             image_files = [
-                Image.open(image) if not isinstance(image, Image.Image) else image
+                stack.enter_context(Image.open(image))
+                if not isinstance(image, Image.Image)
+                else image
                 for image in images
             ]
             assert self.processor is not None, "Processor is not initialized"
@@ -93,15 +95,15 @@ class OnnxImageModel(OnnxModel[T]):
         self,
         model_name: str,
         cache_dir: str,
-        images: Union[ImageInput, Iterable[ImageInput]],
+        images: ImageInput | Iterable[ImageInput],
         batch_size: int = 256,
-        parallel: Optional[int] = None,
-        providers: Optional[Sequence[OnnxProvider]] = None,
-        cuda: bool = False,
-        device_ids: Optional[list[int]] = None,
+        parallel: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
         local_files_only: bool = False,
-        specific_model_path: Optional[str] = None,
-        extra_session_options: Optional[dict[str, Any]] = None,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False

{fastembed_gpu-0.7.4 → fastembed_gpu-0.8.0}/fastembed/image/transform/functional.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 import numpy as np
 from PIL import Image
 
@@ -15,7 +13,7 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 
 def center_crop(
-    image: Union[Image.Image, NumpyArray],
+    image: Image.Image | NumpyArray,
     size: tuple[int, int],
 ) -> NumpyArray:
     if isinstance(image, np.ndarray):
@@ -64,8 +62,8 @@ def center_crop(
 
 def normalize(
     image: NumpyArray,
-    mean: Union[float, list[float]],
-    std: Union[float, list[float]],
+    mean: float | list[float],
+    std: float | list[float],
 ) -> NumpyArray:
     num_channels = image.shape[1] if len(image.shape) == 4 else image.shape[0]
 
@@ -96,8 +94,8 @@ def normalize(
 
 def resize(
     image: Image.Image,
-    size: Union[int, tuple[int, int]],
-    resample: Union[int, Image.Resampling] = Image.Resampling.BILINEAR,
+    size: int | tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.BILINEAR,
 ) -> Image.Image:
     if isinstance(size, tuple):
         return image.resize(size, resample)
@@ -117,7 +115,7 @@ def rescale(image: NumpyArray, scale: float, dtype: type = np.float32) -> NumpyA
     return (image * scale).astype(dtype)
 
 
-def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
+def pil2ndarray(image: Image.Image | NumpyArray) -> NumpyArray:
     if isinstance(image, Image.Image):
         return np.asarray(image).transpose((2, 0, 1))
     return image
@@ -126,7 +124,7 @@ def pil2ndarray(image: Union[Image.Image, NumpyArray]) -> NumpyArray:
 def pad2square(
     image: Image.Image,
     size: int,
-    fill_color: Union[str, int, tuple[int, ...]] = 0,
+    fill_color: str | int | tuple[int, ...] = 0,
 ) -> Image.Image:
     height, width = image.height, image.width
 
@@ -147,3 +145,77 @@ def pad2square(
     new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
     new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
     return new_image
+
+
+def resize_longest_edge(
+    image: Image.Image,
+    max_size: int,
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+) -> Image.Image:
+    height, width = image.height, image.width
+    aspect_ratio = width / height
+
+    if width >= height:
+        # Width is longer
+        new_width = max_size
+        new_height = int(new_width / aspect_ratio)
+    else:
+        # Height is longer
+        new_height = max_size
+        new_width = int(new_height * aspect_ratio)
+
+    # Ensure even dimensions
+    if new_height % 2 != 0:
+        new_height += 1
+    if new_width % 2 != 0:
+        new_width += 1
+
+    return image.resize((new_width, new_height), resample)
+
+
+def crop_ndarray(
+    image: NumpyArray,
+    x1: int,
+    y1: int,
+    x2: int,
+    y2: int,
+    channel_first: bool = True,
+) -> NumpyArray:
+    if channel_first:
+        # (C, H, W) format
+        return image[:, y1:y2, x1:x2]
+    else:
+        # (H, W, C) format
+        return image[y1:y2, x1:x2, :]
+
+
+def resize_ndarray(
+    image: NumpyArray,
+    size: tuple[int, int],
+    resample: int | Image.Resampling = Image.Resampling.LANCZOS,
+    channel_first: bool = True,
+) -> NumpyArray:
+    # Convert to PIL-friendly format (H, W, C)
+    if channel_first:
+        img_hwc = image.transpose((1, 2, 0))
+    else:
+        img_hwc = image
+
+    # Handle different dtypes
+    if img_hwc.dtype == np.float32 or img_hwc.dtype == np.float64:
+        # Assume normalized, scale to 0-255 for PIL
+        img_hwc_scaled = (img_hwc * 255).astype(np.uint8)
+        pil_img = Image.fromarray(img_hwc_scaled, mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized).astype(np.float32) / 255.0
+    else:
+        # uint8 or similar
+        pil_img = Image.fromarray(img_hwc.astype(np.uint8), mode="RGB")
+        resized = pil_img.resize(size, resample)
+        result = np.array(resized)
+
+    # Convert back to original format
+    if channel_first:
+        result = result.transpose((2, 0, 1))
+
+    return result