gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Integration local image embedding to aiox shell

Browse Source
This commit is contained in:
liyaxing
2023-09-26 20:25:02 +08:00
committed by tsukasa
parent 451ab5e0a8
commit 1ff3165961
3 changed files with 44 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/__init__.py
+1 -1
View File
@@ -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
+30 -29
View File
@@ -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,
src/service/aios_shell/aios_shell.py
+13 -13
View File
@@ -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()