Inicial commit of CLIP and GIT simultaneous models

handler.py

@@ -0,0 +1,267 @@
+# handler.py
+import io
+from typing import Any, Dict, List
+
+import numpy as np
+import requests
+import torch
+from PIL import Image
+from transformers import (
+    CLIPModel,
+    CLIPProcessor,
+    CLIPTokenizerFast,
+    pipeline,
+    AutoProcessor,
+    AutoModelForCausalLM,
+)
+from huggingface_hub import logging
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import timeit
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# multi-model list
+multi_model_list = [
+    {"model_id": "openai/clip-vit-base-patch32", "task": "Custom"},
+    {"model_id": "microsoft/git-large-coco", "task": "Custom"},
+]
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        clip_model_id = "openai/clip-vit-base-patch32"
+        # self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.processor_clip = CLIPProcessor.from_pretrained(clip_model_id)
+        self.model_clip = CLIPModel.from_pretrained(clip_model_id).to(self.device)
+        self.tokenizer_clip = CLIPTokenizerFast.from_pretrained(clip_model_id)
+        self.processor_git = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+        self.model_git = AutoModelForCausalLM.from_pretrained(
+            "microsoft/git-large-coco"
+        )
+
+        self.model_git.to(device)
+        self.model_clip.to(device)
+
+        logging.set_verbosity_debug()
+        self.logger = logging.get_logger(__name__)
+
+    def download_image(self, url: str) -> bytes:
+        """
+        Download an image from a given URL.
+
+        Parameters:
+        - url: str
+            The URL from where the image needs to be downloaded.
+
+        Returns:
+        - bytes
+            The downloaded image data in bytes.
+
+        Raises:
+        - Exception: If the image download request fails.
+        """
+        response = requests.get(url)
+        if response.status_code == 200:
+            return response.content
+        else:
+            raise Exception(
+                f"Failed to download image from {url}. Status code: {response.status_code}"
+            )
+
+    def download_images_in_parallel(
+        self, urls: List[str], images_metadata_list: List[dict]
+    ) -> List[bytes]:
+        """
+        Download multiple images in parallel and collect their metadata.
+
+        Parameters:
+        - urls: List[str]
+            A list of URLs from where the images need to be downloaded.
+        - images_metadata_list: List[dict]
+            A list of metadata corresponding to each image URL.
+
+        Returns:
+        - Tuple[List[bytes], List[dict]]
+            A tuple containing a list of downloaded image data in bytes and
+            a list of metadata for the successfully downloaded images.
+        """
+        with ThreadPoolExecutor() as executor:
+            # Start the load operations and mark each future with its URL and metadata
+            future_to_metadata = {
+                executor.submit(self.download_image, url): (url, metadata)
+                for url, metadata in zip(urls, images_metadata_list)
+            }
+
+            results = []
+            successful_metadata = []
+            for future in as_completed(future_to_metadata):
+                url, metadata = future_to_metadata[future]
+                try:
+                    data = future.result()
+                    results.append(data)
+                    metadata["url"] = url
+                    successful_metadata.append(metadata)
+                except Exception as exc:
+                    self.logger.error("%r generated an exception: %s" % (url, exc))
+            return results, successful_metadata
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process the input data based on its type and return the embeddings.
+
+        This method accepts a dictionary with a 'process_type' key that can be either 'images' or 'text'.
+        If 'process_type' is 'images', the method expects a list of image URLs under the 'images_urls' key.
+        It downloads and processes these images, and returns their embeddings.
+        If 'process_type' is 'text', the method expects a string query under the 'query' key.
+        It processes this text and returns its embedding.
+
+        Parameters:
+        - data: Dict[str, Any]
+            A dictionary containing the data to be processed.
+            It must include a 'process_type' key with value either 'images' or 'text'.
+            If 'process_type' is 'images', data should also include 'images_urls' key with a list of image URLs.
+            If 'process_type' is 'text', data should also include 'query' key with a string query.
+
+        Returns:
+        - List[Dict[str, Any]]
+            A list of dictionaries, each containing the embeddings of the processed data.
+            If an error occurs during processing, the dictionary will include an 'error' key with the error message.
+
+        Raises:
+        - ValueError: If the 'process_type' key is not present in data, or if the required keys for 'images' or 'text' are not present or are of the wrong type.
+        """
+
+        if data["process_type"] == "images":
+            try:
+                # Check if 'inputs' key is in data and it has the right type
+                if "images_urls" not in data or not isinstance(
+                    data["images_urls"], list
+                ):
+                    raise ValueError(
+                        "Data must contain 'images_urls' key with a list of images urls."
+                    )
+
+                batch_size = 100
+                if "batch_size" in data:
+                    batch_size = int(data["batch_size"])
+                # Download and process the images (just downloading in this example)
+                images_batches = []
+                processed_metadata = []
+                for i in range(0, len(data["images_urls"]), batch_size):
+                    # select batch of images
+                    batches = data["images_urls"][i : i + batch_size]
+                    batches_metadata = data["images_metadata"][i : i + batch_size]
+
+                    download_start_time = timeit.default_timer()
+
+                    # Download images in parallel along with their metadata
+                    (
+                        downloaded_images,
+                        images_metadata,
+                    ) = self.download_images_in_parallel(batches, batches_metadata)
+
+                    download_end_time = timeit.default_timer()
+                    self.logger.info(
+                        f"Image downloading took {download_end_time - download_start_time} seconds"
+                    )
+                    processing_start_time = timeit.default_timer()
+
+                    for image_content, image_metadata in zip(
+                        downloaded_images, images_metadata
+                    ):
+                        try:
+                            image = Image.open(io.BytesIO(image_content)).convert("RGB")
+                            image_array = np.array(image)
+                            images_batches.append(image_array)
+                            complete_image_metadata = {
+                                "text": image_metadata["caption"],
+                                "source": image_metadata["url"],
+                                "source_type": "images",
+                                **image_metadata,
+                            }
+                            processed_metadata.append(complete_image_metadata)
+
+                        except Exception as e:
+                            print(e)
+                    # This should be a list of images as np.arrays
+                    processing_end_time = timeit.default_timer()
+                    self.logger.info(
+                        f"Image processing took {processing_end_time - processing_start_time} seconds"
+                    )
+
+                embedding_start_time = timeit.default_timer()
+                with torch.no_grad():  # This line ensures that the code inside the block doesn't track gradients
+                    batch = self.processor_clip(
+                        text=None,
+                        images=images_batches,
+                        return_tensors="pt",
+                        padding=True,
+                    )["pixel_values"].to(self.model_clip.device)
+                    batch_git = self.processor_git(
+                        images=images_batches,
+                        return_tensors="pt",
+                    )
+                    git_pixel_values = batch_git.pixel_values.to(self.model_git.device)
+                    # get image captions
+                    generated_ids = self.model_git.generate(
+                        pixel_values=git_pixel_values, max_length=50
+                    )
+
+                    generated_captions = self.processor_git.batch_decode(
+                        generated_ids, skip_special_tokens=True
+                    )
+
+                    # get image embeddings
+                    batch_emb = self.model_clip.get_image_features(pixel_values=batch)
+                    # detach text emb from graph, move to CPU, and convert to numpy array
+                    batch_emb = batch_emb.squeeze(0)
+                    batch_emb = batch_emb.cpu().detach().numpy()
+                    # NORMALIZE
+                    batch_emb = batch_emb.T / np.linalg.norm(batch_emb, axis=1)
+                    # transpose back to (21, 512)
+                    batch_emb = batch_emb.T.tolist()
+                embedding_end_time = timeit.default_timer()
+                self.logger.info(
+                    f"Embedding calculation took {embedding_end_time - embedding_start_time} seconds"
+                )
+
+                # Return the embeddings
+                return {
+                    "embeddings": batch_emb,
+                    "metadata": processed_metadata,
+                    "captions": generated_captions,
+                }
+
+            except Exception as e:
+                print(f"Error during Images processing: {str(e)}")
+                return {"embeddings": [], "error": str(e)}
+
+        elif data["process_type"] == "text":
+            if "query" not in data or not isinstance(data["query"], str):
+                raise ValueError("Data must contain 'query' key which is a str.")
+            query = data["query"]
+            inputs = self.tokenizer_clip(query, return_tensors="pt").to(self.device)
+            text_emb = self.model_clip.get_text_features(**inputs)
+            # detach text emb from graph, move to CPU, and convert to numpy array
+            text_emb = text_emb.detach().cpu().numpy()
+
+            # calculate value to normalize each vector by and normalize them
+            norm_factor = np.linalg.norm(text_emb, axis=1)
+
+            text_emb = text_emb.T / norm_factor
+            # transpose back to (21, 512)
+            text_emb = text_emb.T
+
+            # Converting tensor to list for JSON response
+            text_emb_list = text_emb.tolist()
+
+            return {"embeddings": text_emb_list}
+
+        else:
+            print(
+                f"Error during CLIP endpoint processing: data['process_type']: {data['process_type']} neither 'images' or 'text'"
+            )
+            return {"embeddings": [], "error": str(e)}

requirements.txt

@@ -0,0 +1,24 @@
+certifi==2023.7.22
+charset-normalizer==3.3.2
+colorama==0.4.6
+filelock==3.13.1
+fsspec==2023.10.0
+huggingface-hub==0.17.3
+idna==3.4
+Jinja2==3.1.2
+MarkupSafe==2.1.3
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.24.4
+packaging==23.2
+Pillow==10.1.0
+PyYAML==6.0.1
+regex==2023.10.3
+requests==2.31.0
+safetensors==0.4.0
+sympy==1.12
+tokenizers==0.13.3
+tqdm==4.66.1
+transformers==4.27.2
+typing_extensions==4.8.0
+urllib3==2.0.7