From 85259ca4082504265a5ec099330b2c171eb8fa44 Mon Sep 17 00:00:00 2001 From: Song Date: Sat, 23 Sep 2023 10:44:57 +0800 Subject: [PATCH] Add paint function for agent David --- rootfs/agents/David/agent.toml | 10 +++++- src/aios_kernel/__init__.py | 1 + src/aios_kernel/compute_kernel.py | 32 +++++++++++++++++ src/aios_kernel/compute_task.py | 11 ++++++ src/aios_kernel/local_stability_node.py | 48 ++++++++++++++++++------- src/aios_kernel/workflow_env.py | 22 ++++++++++-- src/service/aios_shell/aios_shell.py | 11 ++++++ test/test_local_sd_node.py | 2 +- 8 files changed, 120 insertions(+), 17 deletions(-) diff --git a/rootfs/agents/David/agent.toml b/rootfs/agents/David/agent.toml index 3462da9..47687f3 100644 --- a/rootfs/agents/David/agent.toml +++ b/rootfs/agents/David/agent.toml @@ -1,6 +1,14 @@ instance_id = "David" fullname = "David" +max_token_size = 16000 +llm_model_name = "gpt-3.5-turbo-16k-0613" [[prompt]] role = "system" -content = "你的名字是David,是一个擅长各种风格的画家。你在收到任何绘画请求时，会尝试产生一组英文单词来说明你想画的画" +content = """你的名字是David,是一个擅长各种风格的画家。你在收到任何绘画请求时，会尝试产生一组英文单词或者短语来描述你想画的画。 +将这些英文单词或者短语作为paint函数的prompt参数。 +当提到风格的时候，需要将风格添加进prompt参数中，而不是作为model_name。 +只有当明确提到用什么模型时，才会去设置model_name参数，否则不要传递model_name参数。 +传递给paint函数的prompt参数一定是一组英文短语，不要用其他语言，如果是其他语言，就将其翻译成英文再传递。 +如果函数调用返回任何错误，直接将错误的原文显示出来并停止。 +如果绘制成功，就停止并输出文件的本地路径。""" diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index 26aadb9..f850184 100644 --- a/src/aios_kernel/__init__.py +++ b/src/aios_kernel/__init__.py @@ -21,6 +21,7 @@ from .storage import ResourceLocation,AIStorage,UserConfig,UserConfigItem from .contact_manager import ContactManager,Contact,FamilyMember from .text_to_speech_function import TextToSpeechFunction from .workspace_env import WorkspaceEnvironment +from .local_stability_node import Local_Stability_ComputeNode AIOS_Version = "0.5.1, build 2023-9-17" diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index fffc8ac..661110b 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -197,3 +197,35 @@ class ComputeKernel: raise Exception("do_text_to_speech failed!") + def text_2_image(self, prompt:str, model_name:Optional[str] = None): + task = ComputeTask() + task.set_text_2_image_params(prompt,model_name) + self.run(task) + return task + + async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None) -> [str, ComputeTaskResult]: + task_req = self.text_2_image(prompt,model_name) + async def check_timer(): + check_times = 0 + while True: + if task_req.state == ComputeTaskState.DONE: + break + + if task_req.state == ComputeTaskState.ERROR: + break + + #may be it will take a long time to generate a image + if check_times >= 60: + task_req.state = ComputeTaskState.ERROR + break + + await asyncio.sleep(0.5) + check_times += 1 + + await asyncio.create_task(check_timer()) + + if task_req.state == ComputeTaskState.DONE: + return None, task_req.result + + return task_req.error_str, None + diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index c6053db..56a8387 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -64,6 +64,17 @@ class ComputeTask: else: self.params["model_name"] = "text-embedding-ada-002" self.params["input"] = input + + def set_text_2_image_params(self, prompt: str, model_name, callchain_id=None): + self.task_type = ComputeTaskType.TEXT_2_IMAGE + self.create_time = time.time() + self.task_id = uuid.uuid4().hex + self.callchain_id = callchain_id + self.params["prompt"] = prompt + if model_name is not None: + self.params["model_name"] = model_name + else: + self.params["model_name"] = "v1-5-pruned-emaonly" def display(self) -> str: return f"ComputeTask: {self.task_id} {self.task_type} {self.state}" diff --git a/src/aios_kernel/local_stability_node.py b/src/aios_kernel/local_stability_node.py index 992e00f..8ee165a 100644 --- a/src/aios_kernel/local_stability_node.py +++ b/src/aios_kernel/local_stability_node.py @@ -6,6 +6,7 @@ import logging import base64 from PIL import Image import requests +from typing import Tuple from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType from .compute_node import ComputeNode @@ -73,6 +74,8 @@ class Local_Stability_ComputeNode(ComputeNode): if self.output_dir is None: self.output_dir = "./" + + self.output_dir = os.path.abspath(self.output_dir) self.start() @@ -85,36 +88,53 @@ class Local_Stability_ComputeNode(ComputeNode): async def remove_task(self, task_id: str): pass + def _make_post_request(self, url, json) -> Tuple[str, requests.Response]: + try: + response = requests.post(url, json=json) + if response.status_code != 200: + return f'{response.status_code}, {response.json()}', None + return None, response + except Exception as e: + return f"{e}", None + + def _run_task(self, task: ComputeTask): task.state = ComputeTaskState.RUNNING model_name = task.params["model_name"] - prompts = task.params["prompts"] + prompt = task.params["prompt"] - logging.info(f"call local stability {model_name} prompts: {prompts}") + logging.info(f"call local stability {model_name} prompts: {prompt}") if model_name is not None: payload = { "sd_model_checkpoint": model_name, } - # {'error': 'RuntimeError', 'detail': '', 'body': '', 'errors': "model 'xxx' not found"} - response = requests.post( - url=f'{self.url}/sdapi/v1/options', json=payload) - if response.status_code != 200: + err, resp = self._make_post_request(f'{self.url}/sdapi/v1/options', payload) + + if err is not None: task.state = ComputeTaskState.ERROR - logger.error( - f"set local stability model failed. err:{response.json()['errors']}") + err_msg = f"Set local stability model failed. err:{err}" + logger.error(err_msg) + task.error_str = err_msg return None + logging.info(f"set local stability model {model_name} success") + payload = { - "prompt": prompts, + "prompt": prompt, "steps": 20 } - response = requests.post( - url=f'{self.url}/sdapi/v1/txt2img', json=payload) - r = response.json() + err, resp = self._make_post_request(f'{self.url}/sdapi/v1/txt2img', payload) + if err is not None: + task.state = ComputeTaskState.ERROR + err_msg = f"Failed. err:{err}" + logger.error(err_msg) + task.error_str = err_msg + return None + + r = resp.json() - print(len(r['images'])) for i in r['images']: image = Image.open(io.BytesIO( base64.b64decode(i.split(",", 1)[0]))) @@ -128,6 +148,8 @@ class Local_Stability_ComputeNode(ComputeNode): return result + task.error_str = "Unknown error!" + task.state = ComputeTaskState.ERROR return None def start(self): diff --git a/src/aios_kernel/workflow_env.py b/src/aios_kernel/workflow_env.py index b7211ef..f872e5b 100644 --- a/src/aios_kernel/workflow_env.py +++ b/src/aios_kernel/workflow_env.py @@ -9,6 +9,7 @@ import logging from typing import Optional from .text_to_speech_function import TextToSpeechFunction +from .compute_kernel import ComputeKernel from .environment import Environment,EnvironmentEvent from .ai_function import SimpleAIFunction from .storage import AIStorage @@ -82,6 +83,16 @@ class CalenderEnvironment(Environment): "update event in calender", self._update_event,update_param)) + + #maybe this function should be in other env? + paint_param = { + "prompt": "A description of the content of the painting", + "model_name": "Which model to use to draw the picture, can be None" + } + self.add_ai_function(SimpleAIFunction("paint", + "Draw a picture according to the description", + self._paint,paint_param)) + #self.add_ai_function(SimpleAIFunction("user_confirm", # "user confirm", # self._user_confirm)) @@ -247,6 +258,14 @@ class CalenderEnvironment(Environment): formatted_time = now.strftime('%Y-%m-%d %H:%M:%S') return formatted_time + + async def _paint(self, prompt, model_name = None) -> str: + err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name) + if err is not None: + return f"exec paint failed. err:{err}" + else: + return f'exec paint OK, saved as a local file, path is: {result.result["file"]}' + # Default Workflow Environment(Context) class WorkflowEnvironment(Environment): def __init__(self, env_id: str,db_file:str) -> None: @@ -327,5 +346,4 @@ class WorkflowEnvironment(Environment): logging.error(f"Error occurred while update env{self.env_id}.{key} ,error:{e}") def get_functions(self): - pass - + pass \ No newline at end of file diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 254996e..a7a003e 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/src/service/aios_shell/aios_shell.py @@ -23,6 +23,7 @@ directory = os.path.dirname(__file__) sys.path.append(directory + '/../../') +from aios_kernel import AIOS_Version,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode,Local_Stability_ComputeNode import proxy from aios_kernel import * @@ -85,6 +86,10 @@ class AIOS_Shell: return False async def initial(self) -> bool: + os.environ["LOCAL_STABILITY_URL"] = "http://192.168.100.79:7866" + os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly" + os.environ["TEXT2IMG_OUTPUT_DIR"] = "./" + cal_env = CalenderEnvironment("calender") await cal_env.start() Environment.set_env_by_id("calender",cal_env) @@ -114,6 +119,12 @@ class AIOS_Shell: await llama_ai_node.start() # ComputeKernel.get_instance().add_compute_node(llama_ai_node) + local_sd_node = Local_Stability_ComputeNode.get_instance() + if await local_sd_node.initial() is not True: + logger.error("local stability node initial failed!") + return False + ComputeKernel.get_instance().add_compute_node(local_sd_node) + await ComputeKernel.get_instance().start() AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg) diff --git a/test/test_local_sd_node.py b/test/test_local_sd_node.py index 62659d1..ebba33b 100644 --- a/test/test_local_sd_node.py +++ b/test/test_local_sd_node.py @@ -30,7 +30,7 @@ async def test_local_sd_node(propmt, model): task.create_time = time.time() task.task_id = uuid.uuid4().hex task.params['model_name'] = model - task.params['prompts'] = propmt + task.params['prompt'] = propmt await node.push_task(task) while True: