Support Intel GPUs using onnxruntime-openvino, run isort

Changes

5 changed files (+47/-10)

• README.md
  • @@ -8,7 +8,7 @@ ```
    pip install -r requirements.txt ``` - For GPU support, install `fastembed-gpu`. + For Nvidia GPUs, install `fastembed-gpu`. For Intel GPUs, install OpenCL drivers (for example `intel-compute-runtime`) and `onnxruntime-openvino`. ## Usage
  • @@ -19,3 +19,9 @@
    Then run `python client.py SEARCH_TEXT NUM_RESULTS` to get a list of the most similar files. You can pass this list to an image viewer such as Gwenview to view image results, although note that Gwenview doesn't preserve the order of the images. Alternatively, add a third parameter to `python client.py` and it will create a temporary directory containing symlinks to the search results and return the directory name. Check out [TIDY](https://github.com/slavabarkov/tidy) (for Android) and [rclip](https://github.com/yurijmikhalevich/rclip) for similar projects, although this one is probably fastest! + + ## TODO + + - Write a `pyproject.toml` and properly package this. + - Support OSes other than Linux? It should be fairly portable now other than a few things like `os.environ["XDG_RUNTIME_DIR"]`. + - GUI
• client.py
  • @@ -2,8 +2,8 @@ import os
    import sys import tempfile import xmlrpc.client - from unixsocket import UnixStreamTransport + from unixsocket import UnixStreamTransport sockpath = os.path.join(os.environ["XDG_RUNTIME_DIR"], "search.sock") proxy = xmlrpc.client.ServerProxy(
• model.py
  • @@ -1,11 +1,42 @@
    + import onnxruntime as ort + import pillow_avif from fastembed import ImageEmbedding, TextEmbedding - import pillow_avif + + print("Loading model") - print("Loading model") + class CustomInferenceSession(ort.InferenceSession): + def __init__( + self, + path_or_bytes, + sess_options=None, + providers=None, + provider_options=None, + **kwargs, + ): + if provider_options is None: + provider_options = [{"device_type": "GPU"}] + super().__init__( + path_or_bytes, + sess_options=sess_options, + providers=providers, + provider_options=provider_options, + **kwargs, + ) + - text_model = TextEmbedding(model_name="jinaai/jina-clip-v1") - image_model = ImageEmbedding(model_name="jinaai/jina-clip-v1") + if "OpenVINOExecutionProvider" in ort.get_available_providers(): + # Use Intel GPU + ort.InferenceSession = CustomInferenceSession + text_model = TextEmbedding( + model_name="jinaai/jina-clip-v1", providers=["OpenVINOExecutionProvider"] + ) + image_model = ImageEmbedding( + model_name="jinaai/jina-clip-v1", providers=["OpenVINOExecutionProvider"] + ) + else: + text_model = TextEmbedding(model_name="jinaai/jina-clip-v1") + image_model = ImageEmbedding(model_name="jinaai/jina-clip-v1") def embed_text(text):
• server.py
  • @@ -4,14 +4,15 @@ import pathlib
    import sqlite3 import sys import threading + import numpy as np import PIL import sqlite_vec + from watchdog.events import FileOpenedEvent, FileSystemEventHandler from watchdog.observers import Observer - from watchdog.events import FileSystemEventHandler, FileOpenedEvent + import model from unixsocket import UnixStreamXMLRPCServer - DIM = 768
• unixsocket.py
  • @@ -1,9 +1,8 @@
    - from http.client import HTTPConnection import socket import socketserver import xmlrpc.client + from http.client import HTTPConnection from xmlrpc.server import SimpleXMLRPCDispatcher, SimpleXMLRPCRequestHandler - # https://stackoverflow.com/questions/11729159/use-python-xmlrpclib-with-unix-domain-sockets