Add local image embedding impls

src/aios_kernel/local_st_compute_node.py

@@ -5,6 +5,7 @@ from pydantic import BaseModel
 
 from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
 from .queue_compute_node import Queue_ComputeNode
+from knowledge import ObjectID
 
 logger = logging.getLogger(__name__)
 
@@ -13,19 +14,20 @@ This is a custom implementation, it should be redesigned.
 """
 
 
-class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode):
+class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
+    # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html
     def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
         super().__init__()
 
-        self.node_id = "local_sentence_transformer_node"
+        self.node_id = "local_sentence_transformer_text_embedding_node"
         self.model_name = model_name
         self.model = None
 
     def initial(self) -> bool:
         logger.info(
-            f"LocalSentenceTransformer_ComputeNode init, model_name: {self.model_name}"
+            f"LocalSentenceTransformer_Text_ComputeNode init, model_name: {self.model_name}"
         )
-        
+
         assert self.model_name is not None
         assert self.model is None
         try:
@@ -35,9 +37,9 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode):
         except Exception as err:
             logger.error(f"load model {self.model} failed: {err}")
             return False
-        
+
         return True
-    
+
     async def execute_task(
         self, task: ComputeTask
     ) -> {
@@ -51,14 +53,14 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode):
         },
     }:
         try:
-            # logger.debug(f"LocalSentenceTransformer_ComputeNode task: {task}")
+            # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}")
             if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
                 input = task.params["input"]
                 logger.debug(
-                    f"LocalSentenceTransformer_ComputeNode task input: {input}"
+                    f"LocalSentenceTransformer_Text_ComputeNode task input: {input}"
                 )
                 sentence_embeddings = self.model.encode(input)
-                # logger.debug(f"LocalSentenceTransformer_ComputeNode task sentence_embeddings: {sentence_embeddings}")
+                # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
                 return {
                     "state": ComputeTaskState.DONE,
                     "content": sentence_embeddings,
@@ -80,9 +82,7 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode):
             }
 
     def display(self) -> str:
-        return (
-            f"LocalSentenceTransformer_ComputeNode: {self.node_id}, {self.model_name}"
-        )
+        return f"LocalSentenceTransformer_Text_ComputeNode: {self.node_id}, {self.model_name}"
 
     def get_capacity(self):
         pass
@@ -94,3 +94,152 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode):
 
     def is_local(self) -> bool:
         return True
+
+from typing import Union
+from PIL import Image
+import io
+            
+class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
+    # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html
+    def __init__(self, model_name: str = "clip-ViT-B-32", multi_model_name: str = "clip-ViT-B-32-multilingual-v1"):
+        super().__init__()
+
+        self.node_id = "local_sentence_transformer_image_embedding_node"
+        self.model_name = model_name
+        self.multi_model_name = multi_model_name
+        self.model = None
+        self.multi_model = None
+
+    def initial(self) -> bool:
+        logger.info(
+            f"LocalSentenceTransformer_Image_ComputeNode init, model_name: {self.model_name} {self.multi_model_name}"
+        )
+
+        assert self.model_name is not None
+        assert self.multi_model_name is not None
+        assert self.model is None
+        assert self.multi_model is None
+        
+        try:
+            from sentence_transformers import SentenceTransformer
+
+            self.model = SentenceTransformer(self.model_name)
+            self.multi_model = SentenceTransformer(self.multi_model_name)
+        except Exception as err:
+            logger.error(f"load model {self.model} failed: {err}")
+            return False
+
+        return True
+
+    def _load_image(self, source: Union[ObjectID, bytes] ) -> Optional[Image]:
+        image_data = None
+        if isinstance(source, ObjectID):
+            from knowledge import KnowledgeStore, ImageObject
+            
+            buf = KnowledgeStore().get_object_store().get_object(source)
+            if buf is None:
+                logger.error(f"load image object but not found! {source}")
+                return None
+            
+            try:
+                image_obj= ImageObject.decode(buf)
+            except Exception as err:
+                logger.error(f"decode ImageObject from buffer failed: {source}, {err}")
+                return None
+            
+            file_size = image_obj.get_file_size()
+            print(f"got image object: {source.to_base58()}, size: {file_size}")
+            
+            image_data = KnowledgeStore().get_chunk_reader().read_chunk_list_to_single_bytes(image_obj.get_chunk_list())
+        
+        elif isinstance(source, bytes):
+            image_data = source
+        else:
+            logger.error(f"unsupport image source type: {type(source)}, {source}")
+            return None
+        
+        try:
+            
+            img = Image.open(io.BytesIO(image_data))
+        
+            return img
+        except Exception as err:
+            logger.error(f"load image from buffer failed: {source}, {err}")
+            return None
+        
+    async def execute_task(
+        self, task: ComputeTask
+    ) -> {
+        "task_type": str,
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        },
+    }:
+        try:
+            # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}")
+            if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
+                input = task.params["input"]
+                logger.debug(
+                    f"LocalSentenceTransformer_Text_ComputeNode task text input: {input}"
+                )
+                sentence_embeddings = self.multi_model.encode(input)
+                # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
+                return {
+                    "state": ComputeTaskState.DONE,
+                    "content": sentence_embeddings,
+                    "message": None,
+                }
+            elif task.task_type == ComputeTaskType.IMAGE_EMBEDDING:
+                input = task.params["input"]
+                logger.debug(
+                    f"LocalSentenceTransformer_Image_ComputeNode task image input: {input}"
+                )
+                
+                img = self._load_image(input)
+                if img is None:
+                    return {
+                        "state": ComputeTaskState.ERROR,
+                        "error": {"code": -1, "message": "load image failed"},
+                    }
+                
+                sentence_embeddings = self.model.encode(img, convert_to_tensor=True)
+                
+                # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
+                return {
+                    "state": ComputeTaskState.DONE,
+                    "content": sentence_embeddings,
+                    "message": None,
+                }
+            else:
+                return {
+                    "state": ComputeTaskState.ERROR,
+                    "error": {"code": -1, "message": "unsupport embedding task type"},
+                }
+        except Exception as err:
+            import traceback
+
+            logger.error(f"{traceback.format_exc()}, error: {err}")
+
+            return {
+                "state": ComputeTaskState.ERROR,
+                "error": {"code": -1, "message": "unknown exception: " + str(err)},
+            }
+
+    def display(self) -> str:
+        return f"LocalSentenceTransformer_Image_ComputeNode: {self.node_id}, {self.model_name}"
+
+    def get_capacity(self):
+        pass
+
+    def is_support(self, task: ComputeTask) -> bool:
+        return (
+            task.task_type == ComputeTaskType.TEXT_EMBEDDING
+            or task.task_type == ComputeTaskType.IMAGE_EMBEDDING
+        )
+
+    def is_local(self) -> bool:
+        return True