diff --git a/c1e8c5afd3d447f2a58727107b5011b7.png b/c1e8c5afd3d447f2a58727107b5011b7.png new file mode 100644 index 0000000..75777f2 Binary files /dev/null and b/c1e8c5afd3d447f2a58727107b5011b7.png differ diff --git a/f7bef6a77c3c4d0087553cf67b4e5f0a.png b/f7bef6a77c3c4d0087553cf67b4e5f0a.png new file mode 100644 index 0000000..f15c96d Binary files /dev/null and b/f7bef6a77c3c4d0087553cf67b4e5f0a.png differ diff --git a/rootfs/agents/David/agent.toml b/rootfs/agents/David/agent.toml index 92cfcf1..0cb02e8 100644 --- a/rootfs/agents/David/agent.toml +++ b/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" diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index f850184..637ee7d 100644 --- a/src/aios_kernel/__init__.py +++ b/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" diff --git a/src/aios_kernel/stability_node.py b/src/aios_kernel/stability_node.py index b36e86c..88eeb71 100644 --- a/src/aios_kernel/stability_node.py +++ b/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: - self.model = os.getenv("STABILITY_MODEL") + if os.getenv("STABILITY_DEFAULT_MODEL") is not None: + 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( - key=self.api_key, - verbose=True, # Print debug messages. - engine=self.model, - ) + if os.getenv("TEXT2IMG_OUTPUT_DIR") is not None: + self.output_dir = os.getenv("TEXT2IMG_OUTPUT_DIR") + else: + 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 - prompts = task.params["prompts"] + model_name = task.params["model_name"] + prompt = task.params["prompt"] - logging.info(f"call stability {self.model} prompts: {prompts}") - answers = self.client.generate( - prompt=prompts, + logging.info(f"call stability {self.default_model} prompts: {prompt}") + + 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): diff --git a/src/aios_kernel/workflow_env.py b/src/aios_kernel/workflow_env.py index 39dda7b..f2b41d3 100644 --- a/src/aios_kernel/workflow_env.py +++ b/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: diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 3cde4b6..c525b47 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/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: diff --git a/test/test_local_sd_node.py b/test/test_local_sd_node.py index ebba33b..949cc7d 100644 --- a/test/test_local_sd_node.py +++ b/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"] = "http://192.168.100.79:7860" os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly" os.environ["TEXT2IMG_OUTPUT_DIR"] = "./" diff --git a/test/test_sd_api.py b/test/test_sd_api.py new file mode 100644 index 0000000..0f4544a --- /dev/null +++ b/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))