Integration local image embedding to aiox shell

src/aios_kernel/__init__.py

@@ -23,7 +23,7 @@ from .text_to_speech_function import TextToSpeechFunction
 from .workspace_env import WorkspaceEnvironment
 from .local_stability_node import Local_Stability_ComputeNode
 from .stability_node import Stability_ComputeNode
-from .local_st_compute_node import LocalSentenceTransformer_ComputeNode
+from .local_st_compute_node import LocalSentenceTransformer_Text_ComputeNode,LocalSentenceTransformer_Image_ComputeNode
 
 AIOS_Version = "0.5.1, build 2023-9-26"
 

src/aios_kernel/local_st_compute_node.py

@@ -2,16 +2,15 @@ import logging
 import requests
 from typing import Optional, List
 from pydantic import BaseModel
-
+from typing import Union
+from PIL import Image
+import io
 from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
 from .queue_compute_node import Queue_ComputeNode
 from knowledge import ObjectID
 
 logger = logging.getLogger(__name__)
 
-"""
-This is a custom implementation, it should be redesigned.
-"""
 
 
 class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
@@ -88,20 +87,19 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
         pass
 
     def is_support(self, task: ComputeTask) -> bool:
-        return task.task_type == ComputeTaskType.TEXT_EMBEDDING and (
+        return task.task_type == ComputeTaskType.TEXT_EMBEDDING
-            not task.params["model_name"] or task.params["model_name"] == "llama"
-        )
 
     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"):
+    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"
@@ -119,7 +117,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
         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
 
@@ -131,42 +129,45 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
 
         return True
 
-    def _load_image(self, source: Union[ObjectID, bytes] ) -> Optional[Image]:
+    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)
+                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())
+            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
     ) -> {
@@ -198,16 +199,16 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
                 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)
+                sentence_embeddings = self.model.encode(img)
-                
+
                 # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
                 return {
                     "state": ComputeTaskState.DONE,

src/service/aios_shell/aios_shell.py

@@ -147,6 +147,7 @@ class AIOS_Shell:
                 logger.error("llama node initial failed!")
                 await AIStorage.get_instance().set_feature_init_result("llama",False)
 
+
         if await AIStorage.get_instance().is_feature_enable("aigc"):
             try:
                 google_text_to_speech_node = GoogleTextToSpeechNode.get_instance()
@@ -161,21 +162,20 @@ class AIOS_Shell:
             #     logger.error("stability api node initial failed!")
             # ComputeKernel.get_instance().add_compute_node(stability_api_node)
 
-        local_sd_node = Local_Stability_ComputeNode.get_instance()
+        
-        if await local_sd_node.initial() is True:
+        
-            ComputeKernel.get_instance().add_compute_node(local_sd_node)
+        local_st_text_compute_node = LocalSentenceTransformer_Text_ComputeNode()
+        if local_st_text_compute_node.initial() is not True:
+            logger.error("local sentence transformer text embedding node initial failed!")
         else:
-            logger.error("local stability node initial failed!")
+            ComputeKernel.get_instance().add_compute_node(local_st_text_compute_node)
-            await AIStorage.get_instance.set_feature_init_result("aigc",False)
+            
-
+        local_st_image_compute_node = LocalSentenceTransformer_Image_ComputeNode()
-
+        if local_st_image_compute_node.initial() is not True:
-    
+            logger.error("local sentence transformer image embedding node initial failed!")
-        local_st_compute_node = LocalSentenceTransformer_ComputeNode()
-        if local_st_compute_node.initial() is not True:
-            logger.error("local sentence transformer node initial failed!")
         else:
-        	ComputeKernel.get_instance().add_compute_node(local_st_compute_node)
+            ComputeKernel.get_instance().add_compute_node(local_st_image_compute_node)
-
+       
 
         await ComputeKernel.get_instance().start()