Merge pull request #66 from glen0125/MVP

Add paint env & test stability api

rootfs/agents/David/agent.toml

@@ -2,7 +2,7 @@ instance_id = "David"
 fullname = "David"
 max_token_size = 16000
 llm_model_name = "gpt-3.5-turbo-16k-0613"
-owner_env = "calender"
+owner_env = "paint"
 
 [[prompt]]
 role = "system"

src/aios_kernel/__init__.py

@@ -10,7 +10,7 @@ from .knowledge_pipeline import KnowledgeEmailSource, KnowledgeDirSource, Knowle
 from .role import AIRole,AIRoleGroup
 from .workflow import Workflow
 from .bus import AIBus
-from .workflow_env import WorkflowEnvironment,CalenderEnvironment,CalenderEvent
+from .workflow_env import WorkflowEnvironment,CalenderEnvironment,CalenderEvent,PaintEnvironment
 from .local_llama_compute_node import LocalLlama_ComputeNode
 from .whisper_node import WhisperComputeNode
 from .google_text_to_speech_node import GoogleTextToSpeechNode
@@ -22,6 +22,7 @@ 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
+from .stability_node import Stability_ComputeNode
 
 AIOS_Version = "0.5.1, build 2023-9-17"
 

src/aios_kernel/stability_node.py

@@ -3,6 +3,7 @@ import io
 import asyncio
 from asyncio import Queue
 import logging
+from pathlib import Path
 
 from PIL import Image
 from stability_sdk import client
@@ -16,7 +17,7 @@ logger = logging.getLogger(__name__)
 
 
 class Stability_ComputeNode(ComputeNode):
-    _instanace = None
+    _instance = None
 
     @classmethod
     def get_instance(cls):
@@ -31,6 +32,15 @@ class Stability_ComputeNode(ComputeNode):
             "stability_api_key", "stability api key", False, None)
         user_config.add_user_config(
             "stability_model", "stability model name", True, "stable-diffusion-512-v2-1")
+        if os.getenv("TEXT2IMG_OUTPUT_DIR") is None:
+            home_dir = Path.home()
+            output_dir = Path.joinpath(home_dir, "text2img_output")
+            Path.mkdir(output_dir, exist_ok=True)
+            user_config.add_user_config(
+                "text2img_output_dir", "text2image output dir", True, output_dir)
+        if os.getenv("STABILITY_DEFAULT_MODEL") is None:
+            user_config.add_user_config(
+                "stability_default_model", "stability default model", True, "stable-diffusion-512-v2-1")
 
     def __init__(self):
         super().__init__()
@@ -38,10 +48,11 @@ class Stability_ComputeNode(ComputeNode):
         self.is_start = False
         self.node_id = "stability_node"
         self.api_key = ""
-        self.model = ""
+        self.default_model = ""
 
         self.task_queue = Queue()
 
+    async def initial(self):
         if os.getenv("STABILITY_API_KEY") is not None:
             self.api_key = os.getenv("STABILITY_API_KEY")
         else:
@@ -53,16 +64,23 @@ class Stability_ComputeNode(ComputeNode):
             return False
 
         # Check out the following link for a list of available engines: https://platform.stability.ai/docs/features/api-parameters#engine
-        if os.getenv("STABILITY_MODEL") is not None:
+        if os.getenv("STABILITY_DEFAULT_MODEL") is not None:
-            self.model = os.getenv("STABILITY_MODEL")
+            self.default_model = os.getenv("STABILITY_DEFAULT_MODEL")
         else:
-            self.model = AIStorage.get_instance().get_user_config().get_value("stability_model")
+            self.default_model = AIStorage.get_instance().get_user_config().get_value("stability_default_model")
+        
+        if self.default_model is None:
+            self.default_model = "stable-diffusion-512-v2-1"
 
-        self.client = client.StabilityInference(
+        if os.getenv("TEXT2IMG_OUTPUT_DIR") is not None:
-            key=self.api_key,
+            self.output_dir = os.getenv("TEXT2IMG_OUTPUT_DIR")
-            verbose=True,  # Print debug messages.
+        else:
-            engine=self.model,
+            self.output_dir = AIStorage.get_instance(
-        )
+            ).get_user_config().get_value("text2img_output_dir")
+
+        if self.output_dir is None:
+            self.output_dir = "./"
+            self.output_dir = os.path.abspath(self.output_dir)
 
         self.start()
 
@@ -77,12 +95,26 @@ class Stability_ComputeNode(ComputeNode):
 
     def _run_task(self, task: ComputeTask):
         task.state = ComputeTaskState.RUNNING
-        # model_name && max_token_size not used here
+        model_name = task.params["model_name"]
-        prompts = task.params["prompts"]
+        prompt = task.params["prompt"]
 
-        logging.info(f"call stability {self.model} prompts: {prompts}")
+        logging.info(f"call stability {self.default_model} prompts: {prompt}")
-        answers = self.client.generate(
+
-            prompt=prompts,
+        api = None
+        try:
+            api = client.StabilityInference(
+                key=self.api_key,
+                verbose=True,  # Print debug messages.
+                engine=model_name,
+            )
+        except Exception as e:
+            task.error_str = f"create stability client failed: {e}"
+            logging.warn(task.error_str)
+            task.state = ComputeTaskState.ERROR
+            return None
+
+        answers = api.generate(
+            prompt=prompt,
             # If a seed is provided, the resulting generated image will be deterministic.
             seed=0,
             # What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again.
@@ -105,15 +137,16 @@ class Stability_ComputeNode(ComputeNode):
 
         for resp in answers:
             for artifact in resp.artifacts:
-                logger.info(
-                    f"artifact:{artifact.id},{artifact.type},{artifact.finish_reason}")
-
                 if artifact.finish_reason == generation.FILTER:
-                    logging.warn("request activated the API's safety filters")
+                    err_msg = "request activated the API's safety filters"
+                    logging.warn(err_msg)
+                    task.error_str = err_msg
+                    task.state = ComputeTaskState.ERROR
+                    return None
                 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 = task.task_id + ".png"  # which dir to save?
+                    file_name = os.path.join(self.output_dir, task.task_id + ".png")
                     img.save(file_name)
 
                     result = ComputeTaskResult()
@@ -123,6 +156,8 @@ class 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

@@ -317,6 +317,32 @@ class CalenderEnvironment(Environment):
         else:
             return f'exec paint OK, saved as a local file, path is: {result.result["file"]}'
 
+
+class PaintEnvironment(Environment):
+    def __init__(self, env_id: str) -> None:
+        super().__init__(env_id)
+        self.is_run = False
+
+        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))
+
+    def _do_get_value(self,key:str) -> Optional[str]:
+        return None
+
+
+    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:

src/service/aios_shell/aios_shell.py

@@ -24,7 +24,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
+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,Stability_ComputeNode,PaintEnvironment
 from aios_kernel import ContactManager,Contact
 import proxy
 from aios_kernel import *
@@ -77,7 +77,7 @@ class AIOS_Shell:
 
         Local_Stability_ComputeNode.declare_user_config()
 
-        
+        #Stability_ComputeNode.declare_user_config()
 
 
 
@@ -121,6 +121,9 @@ class AIOS_Shell:
         workspace_env = WorkspaceEnvironment("bash")
         Environment.set_env_by_id("bash",workspace_env)
 
+        paint_env = PaintEnvironment("paint")
+        Environment.set_env_by_id("paint",paint_env)
+
         if await AgentManager.get_instance().initial() is not True:
             logger.error("agent manager initial failed!")
             return False
@@ -152,6 +155,10 @@ class AIOS_Shell:
                 logger.error(f"google text to speech node initial failed! {e}")
                 await AIStorage.get_instance.set_feature_init_result("aigc",False)
 
+            # stability_api_node = Stability_ComputeNode()
+            # if await stability_api_node.initial() is not True:
+            #     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:

test/test_local_sd_node.py

@@ -14,7 +14,7 @@ from aios_kernel.compute_task import ComputeTaskType, ComputeTask, ComputeTaskSt
 #to launch a local stability node, please check:
 #https://github.com/glen0125/stable-diffusion-webui-docker
 
-os.environ["LOCAL_STABILITY_URL"] = "http://aigc:7866"
+os.environ["LOCAL_STABILITY_URL"] = ""
 os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly"
 os.environ["TEXT2IMG_OUTPUT_DIR"] = "./"
 

test/test_sd_api.py

@@ -0,0 +1,56 @@
+import os
+import time
+import uuid
+import io
+import asyncio
+import sys
+import logging
+import pytest
+directory = os.path.dirname(__file__)
+sys.path.append(directory + '/../src')
+from aios_kernel.stability_node import Stability_ComputeNode
+from aios_kernel.compute_task import ComputeTaskType, ComputeTask, ComputeTaskState
+
+os.environ["STABILITY_API_KEY"] = ""
+os.environ["STABILITY_DEFAULT_MODEL"] = "stable-diffusion-512-v2-1"
+os.environ["TEXT2IMG_OUTPUT_DIR"] = "./"
+
+@pytest.mark.asyncio
+async def test_sd_api(propmt, model):
+    node = Stability_ComputeNode.get_instance()
+    if await node.initial() is not True:
+        print("node initial failed!")
+        return
+    
+    task = ComputeTask()
+    task.task_type = ComputeTaskType.TEXT_2_IMAGE
+    task.create_time = time.time()
+    task.task_id = uuid.uuid4().hex
+    task.params['model_name'] = model
+    task.params['prompt'] = propmt
+    await node.push_task(task)
+
+    while True:
+        if task.state == ComputeTaskState.DONE:
+            local_file = task.result.result
+            print("local file is: ", local_file)
+            break
+        await asyncio.sleep(1)
+
+     # result = node._run_task(task)
+    # print("result is: ", result)
+
+    # if result.result is not None:
+    #     local_file = result.result
+    #     print("local file is: ", local_file)
+
+if __name__ == "__main__":
+    arg_len = len(os.sys.argv)
+    prompt = "a beautiful sunset"
+    model = "stable-diffusion-512-v2-1"
+    if arg_len >= 2:
+        prompt = os.sys.argv[1]
+    if arg_len == 3:
+        model = os.sys.argv[2]
+
+    asyncio.run(test_sd_api(prompt, model))