Merge pull request #67 from glen0125/MVP

Return result code from sd node

rootfs/agents/David/agent.toml

@@ -1,15 +1,14 @@
 instance_id = "David"
 fullname = "David"
 max_token_size = 16000
-llm_model_name = "gpt-3.5-turbo-16k-0613"
+
 owner_env = "paint"
 
 [[prompt]]
 role = "system"
-content = """你的名字是David,是一个擅长各种风格的画家。你在收到任何绘画请求时，会尝试产生一组英文单词或者短语来描述你想画的画。
+content = """You are an artist, and you will use the 'paint' function in the llm_inner_functions to create artwork. 
-将这些英文单词或者短语作为paint函数的prompt参数。
+You will extract relevant keywords based on user input and requirements, and pass these keywords to the 'prompt' parameter of the 'paint' function. 
-当提到风格的时候，需要将风格添加进prompt参数中，而不是作为model_name。
+When a specific model is mentioned, pass the model's name to the 'model_name' parameter of the 'paint' function. 
-只有当明确提到用什么模型时，才会去设置model_name参数，否则不要传递model_name参数。
+If there are instructions not to depict certain content, summarize this content and pass it as keywords to the 'negative_prompt' parameter. 
-传递给paint函数的prompt参数一定是一组英文短语，不要用其他语言，如果是其他语言，就将其翻译成英文再传递。
+All parameters must be in English. 
-如果函数调用返回任何错误，直接将错误的原文显示出来并停止。
+When the user mentions creating a children's picture book, use the "realisticVisionV51_v51VAE" model, and append the keyword "<lora:COOLKIDS_MERGE_V2.5:1> <lora:add_detail:-0.5>" to the 'prompt' parameter."""
-如果绘制成功，就停止并输出文件的本地路径。"""

src/aios_kernel/compute_kernel.py

@@ -175,17 +175,19 @@ class ComputeKernel:
             return task_result.result
 
 
-    def text_2_image(self, prompt:str, model_name:Optional[str] = None):
+    def text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None):
         task = ComputeTask()
-        task.set_text_2_image_params(prompt,model_name)
+        task.set_text_2_image_params(prompt,model_name, negative_prompt)
         self.run(task)
         return task
 
-    async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None) -> [str, ComputeTaskResult]:
+    async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult:
-        task_req = self.text_2_image(prompt,model_name)
+        task = self.text_2_image(prompt,model_name, negative_prompt)
-        task_result = self._send_task(task_req)
+        task = await self._send_task(task)
-        if task_req.state == ComputeTaskState.DONE:
-            return None, task_result
 
-        return task_req.error_str, None
+        return task.result
+        # if task_req.state == ComputeTaskState.DONE:
+        #     return None, task_result
+
+        # return task_req.error_str, None
 

src/aios_kernel/compute_task.py

@@ -71,12 +71,13 @@ class ComputeTask:
             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):
+    def set_text_2_image_params(self, prompt: str, model_name, negative_prompt="", 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
+        self.params["negative_prompt"] = negative_prompt
         if model_name is not None:
             self.params["model_name"] = model_name
         else:

src/aios_kernel/local_stability_node.py

@@ -9,7 +9,7 @@ import requests
 from typing import Tuple
 from pathlib import Path
 
-from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType, ComputeTaskResultCode
 from .compute_node import ComputeNode
 from .storage import AIStorage, UserConfig
 
@@ -107,10 +107,19 @@ class Local_Stability_ComputeNode(ComputeNode):
 
     def _run_task(self, task: ComputeTask):
         task.state = ComputeTaskState.RUNNING
+        result = ComputeTaskResult()
+        result.result_code = ComputeTaskResultCode.ERROR
+        result.set_from_task(task)
+        
         model_name = task.params["model_name"]
         prompt = task.params["prompt"]
+        negative_prompt = task.params["negative_prompt"]
+        if negative_prompt == None or negative_prompt == "":
+            negative_prompt = "sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, duplicate, mutated hands, mutated legs, (blurry:1.3), (bad anatomy:1.2), bad proportions, extra limbs, more than 2 nipples, extra legs, fused fingers, missing fingers, jpeg artifacts, signature, watermark, username, artist name, heterochromia, muscular legs, monochrome, grayscale, skin spots, acnes, skin blemishes, age spot, skin spots, acnes, logo, badhandv4, easynegative, cropped image, patreon,lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, ng_deepnegative_v1_75t, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry,(Tiptoe:1.3),looking at viewer, Twisted eyes"
 
-        logging.info(f"call local stability {model_name} prompts: {prompt}")
+        prompt += ",masterpiece, best quality:1.3"
+
+        logging.info(f"call local stability {model_name} prompts: {prompt}, nagative_prompt: {negative_prompt}")
 
         if model_name is not None:
             payload = {
@@ -123,12 +132,14 @@ class Local_Stability_ComputeNode(ComputeNode):
                 err_msg = f"Set local stability model failed. err:{err}"
                 logger.error(err_msg)
                 task.error_str = err_msg
-                return None
+                result.error_str = err_msg
+                return result
 
             logging.info(f"set local stability model {model_name} success")
 
         payload = {
             "prompt": prompt,
+            "negative_prompt": negative_prompt,
             "steps": 20
         }
 
@@ -138,7 +149,8 @@ class Local_Stability_ComputeNode(ComputeNode):
             err_msg = f"Failed. err:{err}"
             logger.error(err_msg)
             task.error_str = err_msg
-            return None
+            result.error_str = err_msg
+            return result
         
         r = resp.json()
 
@@ -148,16 +160,17 @@ class Local_Stability_ComputeNode(ComputeNode):
             file_name = os.path.join(self.output_dir, task.task_id + ".png")
             image.save(file_name)
 
-            result = ComputeTaskResult()
+            task.state = ComputeTaskState.DONE
-            result.set_from_task(task)
+            result.result_code = ComputeTaskResultCode.OK
             result.worker_id = self.node_id
             result.result = {"file": file_name}
 
             return result
 
         task.error_str = "Unknown error!"
+        result.error_str = "Unknown error!"
         task.state = ComputeTaskState.ERROR
-        return None
+        return result
 
     def start(self):
         if self.is_start:
@@ -170,9 +183,9 @@ class Local_Stability_ComputeNode(ComputeNode):
                 task = await self.task_queue.get()
                 logger.info(f"stability_node get task: {task.display()}")
                 result = self._run_task(task)
-                if result is not None:
+                # if result is not None:
-                    task.state = ComputeTaskState.DONE
+                #     task.state = ComputeTaskState.DONE
-                    task.result = result
+                #     task.result = result
 
         asyncio.create_task(_run_task_loop())
 

src/aios_kernel/stability_node.py

@@ -9,7 +9,7 @@ from PIL import Image
 from stability_sdk import client
 import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
 
-from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType, ComputeTaskResultCode
 from .compute_node import ComputeNode
 from .storage import AIStorage, UserConfig
 
@@ -95,10 +95,15 @@ class Stability_ComputeNode(ComputeNode):
 
     def _run_task(self, task: ComputeTask):
         task.state = ComputeTaskState.RUNNING
+        result = ComputeTaskResult()
+        result.result_code = ComputeTaskResultCode.ERROR
+        result.set_from_task(task)
+        
         model_name = task.params["model_name"]
         prompt = task.params["prompt"]
+        negative_prompt = task.params["negative_prompt"]
 
-        logging.info(f"call stability {self.default_model} prompts: {prompt}")
+        logging.info(f"call stability {self.default_model} prompts: {prompt}, negative_prompt: {negative_prompt}")
 
         api = None
         try:
@@ -109,9 +114,10 @@ class Stability_ComputeNode(ComputeNode):
             )
         except Exception as e:
             task.error_str = f"create stability client failed: {e}"
+            result.error_str = f"create stability client failed: {e}"
             logging.warn(task.error_str)
             task.state = ComputeTaskState.ERROR
-            return None
+            return result
 
         answers = api.generate(
             prompt=prompt,
@@ -141,24 +147,26 @@ class Stability_ComputeNode(ComputeNode):
                     err_msg = "request activated the API's safety filters"
                     logging.warn(err_msg)
                     task.error_str = err_msg
+                    result.error_str = err_msg
                     task.state = ComputeTaskState.ERROR
-                    return None
+                    return result
                 if artifact.type == generation.ARTIFACT_IMAGE:
                     img = Image.open(io.BytesIO(artifact.binary))
                     # Save our generated images with the task_id as the filename.
                     file_name = os.path.join(self.output_dir, task.task_id + ".png")
                     img.save(file_name)
 
-                    result = ComputeTaskResult()
+                    task.state = ComputeTaskState.DONE
-                    result.set_from_task(task)
+                    result.result_code = ComputeTaskResultCode.OK
                     result.worker_id = self.node_id
                     result.result = {"file": file_name}
 
                     return result
 
         task.error_str = "Unknown error!"
+        result.error_str = "Unknown error!"
         task.state = ComputeTaskState.ERROR
-        return None
+        return result
 
     def start(self):
         if self.is_start:
@@ -171,9 +179,9 @@ class Stability_ComputeNode(ComputeNode):
                 task = await self.task_queue.get()
                 logger.info(f"stability_node get task: {task.display()}")
                 result = self._run_task(task)
-                if result is not None:
+                # if result is not None:
-                    task.state = ComputeTaskState.DONE
+                #     task.state = ComputeTaskState.DONE
-                    task.result = result
+                #     task.result = result
 
         asyncio.create_task(_run_task_loop())
 

src/aios_kernel/workflow_env.py

@@ -9,7 +9,7 @@ import logging
 from typing import Optional
 
 from .text_to_speech_function import TextToSpeechFunction
-from .compute_kernel import ComputeKernel
+from .compute_kernel import ComputeKernel, ComputeTaskResultCode
 from .environment import Environment,EnvironmentEvent
 from .ai_function import SimpleAIFunction
 from .storage import AIStorage
@@ -311,9 +311,9 @@ class CalenderEnvironment(Environment):
 
     
     async def _paint(self, prompt, model_name = None) -> str:
-        err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name)
+        result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name)
-        if err is not None:
+        if result.result_code == ComputeTaskResultCode.ERROR:
-            return f"exec paint failed. err:{err}"
+            return f"exec paint failed. err:{result.error_str}"
         else:
             return f'exec paint OK, saved as a local file, path is: {result.result["file"]}'
 
@@ -324,19 +324,20 @@ class PaintEnvironment(Environment):
         self.is_run = False
 
         paint_param = {
-            "prompt": "A description of the content of the painting",
+            "prompt": "Keywords of the content of the painting",
-            "model_name": "Which model to use to draw the picture, can be None"
+            "model_name": "Which model to use to draw the picture, can be None",
+            "negative_prompt": "Keywords that describe what is not to be drawn, can be None"
         }
         self.add_ai_function(SimpleAIFunction("paint",
-                                        "Draw a picture according to the description",
+                                        "Draw a picture according to the keywords",
                                         self._paint,paint_param))
 
     def _do_get_value(self,key:str) -> Optional[str]:
         return None
 
 
-    async def _paint(self, prompt, model_name = None) -> str:
+    async def _paint(self, prompt, model_name = None, negative_prompt = None) -> str:
-        err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name)
+        err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name, negative_prompt)
         if err is not None:
             return f"exec paint failed. err:{err}"
         else: