diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index 2be558b..ef57ea9 100644 --- a/src/aios_kernel/__init__.py +++ b/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" diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py index 87fe37d..2f36936 100644 --- a/src/aios_kernel/local_st_compute_node.py +++ b/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" @@ -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, diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 736be84..4ba88b6 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/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!") - await AIStorage.get_instance.set_feature_init_result("aigc",False) - - - - local_st_compute_node = LocalSentenceTransformer_ComputeNode() - if local_st_compute_node.initial() is not True: - logger.error("local sentence transformer node initial failed!") + ComputeKernel.get_instance().add_compute_node(local_st_text_compute_node) + + 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()