Add paint function for agent David

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

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