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 (
-            not task.params["model_name"] or task.params["model_name"] == "llama"
-        )
+        return task.task_type == ComputeTaskType.TEXT_EMBEDDING
 
     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"
@@ -150,7 +148,11 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
             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
@@ -159,7 +161,6 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
             return None
 
         try:
-            
             img = Image.open(io.BytesIO(image_data))
 
             return img
@@ -206,7 +207,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
                         "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 {

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,20 +162,19 @@ 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!")
-            await AIStorage.get_instance.set_feature_init_result("aigc",False)
+            ComputeKernel.get_instance().add_compute_node(local_st_text_compute_node)
             
-
-    
-        local_st_compute_node = LocalSentenceTransformer_ComputeNode()
-        if local_st_compute_node.initial() is not True:
-            logger.error("local sentence transformer node initial failed!")
+        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!")
         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()