Add paint function for agent David

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参数一定是一组英文短语，不要用其他语言，如果是其他语言，就将其翻译成英文再传递。
+如果函数调用返回任何错误，直接将错误的原文显示出来并停止。
+如果绘制成功，就停止并输出文件的本地路径。"""

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"
 

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
+

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}"

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"}
+            err, resp = self._make_post_request(f'{self.url}/sdapi/v1/options', payload)
-            response = requests.post(
+
-                url=f'{self.url}/sdapi/v1/options', json=payload)
+            if err is not None:
-            if response.status_code != 200:
                 task.state = ComputeTaskState.ERROR
-                logger.error(
+                err_msg = f"Set local stability model failed. err:{err}"
-                    f"set local stability model failed. err:{response.json()['errors']}")
+                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(
+        err, resp = self._make_post_request(f'{self.url}/sdapi/v1/txt2img', payload)
-            url=f'{self.url}/sdapi/v1/txt2img', json=payload)
+        if err is not None:
-        r = response.json()
+            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):

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
-

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)

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: