diff --git a/rootfs/agents/David/agent.toml b/rootfs/agents/David/agent.toml index 47d2ff8..df8c308 100644 --- a/rootfs/agents/David/agent.toml +++ b/rootfs/agents/David/agent.toml @@ -6,9 +6,17 @@ owner_env = "paint" [[prompt]] role = "system" -content = """You are an artist, and you will use the 'paint' function in the llm_inner_functions to create artwork. -You will extract relevant keywords based on user input and requirements, and pass these keywords to the 'prompt' parameter of the 'paint' function. -When a specific model is mentioned, pass the model's name to the 'model_name' parameter of the 'paint' function. -If there are instructions not to depict certain content, summarize this content and pass it as keywords to the 'negative_prompt' parameter. -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 " " to the 'prompt' parameter.""" +content = """You are an advanced artist AI, and your task is to understand users' textual descriptions and transform these descriptions into nuanced visual images. When a user provides a description, you need to carefully analyze and comprehend the details, emotions, environment, objects, and style within the description. Then, you will use this information to create an image that accurately and creatively reflects the user's description. + +Your creative process is as follows: First, you will receive a description from the user, which could be a depiction of anything, such as a scene, an object, a concept, or an emotional expression. You need to extract key information from these descriptions, such as the subject, atmosphere, color scheme, time period, and artistic style. Then, you will pass these key pieces of information as parameters to the prompt parameter of your "paint" function. + +For example, if a user says, "I want a painting that depicts a sunset on a summer beach with children playing and sailboats in the distance, with the sky showing a gradient of orange and purple." You should extract: +- Time: Summer, sunset moment +- Scene: Beach +- Activity: Children playing +- Objects: Sailboats in the distance +- Colors: The gradient of orange and purple in the sky + +Then you will combine these pieces of information into an accurate prompt to pass to the "paint" function, for example: "During a summer sunset, children play on the beach with sailboats visible in the distance, and the sky displays a gradient of orange and purple." + +Please always ensure that your creations respect the user's intentions and that you infuse your unique artistic style and creativity into your work.""" diff --git a/src/aios/environment/dall_e_compute_node.py b/src/aios/environment/dall_e_compute_node.py new file mode 100644 index 0000000..8e63e03 --- /dev/null +++ b/src/aios/environment/dall_e_compute_node.py @@ -0,0 +1,146 @@ +import os +import io +import asyncio +from asyncio import Queue +import logging +from pathlib import Path +from openai import OpenAI +import base64 + +from PIL import Image + +from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType, ComputeTaskResultCode +from .compute_node import ComputeNode +from .storage import AIStorage, UserConfig + +logger = logging.getLogger(__name__) + + +class DallE_ComputeNode(ComputeNode): + _instance = None + + @classmethod + def get_instance(cls): + if cls._instance is None: + cls._instance = DallE_ComputeNode() + return cls._instance + + @classmethod + def declare_user_config(cls): + user_config = AIStorage.get_instance().get_user_config() + + 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) + + def __init__(self): + super().__init__() + + self.is_start = False + self.node_id = "dall_e_node" + self.openai_api_key = "" + self.default_model = "dall-e-3" + + self.task_queue = Queue() + + async def initial(self): + if os.getenv("OPENAI_API_KEY") is not None: + self.openai_api_key = os.getenv("OPENAI_API_KEY") + else: + self.openai_api_key = AIStorage.get_instance().get_user_config().get_value("openai_api_key") + + if self.openai_api_key is None: + logger.error("openai_api_key is None!") + return False + + 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() + + return True + + async def push_task(self, task: ComputeTask, proiority: int = 0): + logger.info(f"DallE_node push task: {task.display()}") + self.task_queue.put_nowait(task) + + async def remove_task(self, task_id: str): + pass + + def _run_task(self, task: ComputeTask): + task.state = ComputeTaskState.RUNNING + result = ComputeTaskResult() + result.result_code = ComputeTaskResultCode.ERROR + result.set_from_task(task) + + try: + prompt = task.params["prompt"] + logging.info(f"Call DallE {self.default_model} prompts: {prompt}") + client = OpenAI(api_key=self.openai_api_key) + + response = client.images.generate( + model=self.default_model, + prompt=prompt, + size="1024x1024", + quality="standard", + n=1, + response_format="b64_json", + ) + + binary_data = base64.b64decode(response.data[0].b64_json) + image = Image.open(io.BytesIO(binary_data)) + file_name = os.path.join(self.output_dir, task.task_id + ".png") + image.save(file_name) + + task.state = ComputeTaskState.DONE + result.result_code = ComputeTaskResultCode.OK + result.worker_id = self.node_id + result.result = {"file": file_name} + + return result + + except Exception as e: + logging.error(f"Call DallE failed. err: {e}") + task.error_str = str(e) + result.error_str = str(e) + task.state = ComputeTaskState.ERROR + return result + + async def start(self): + if self.is_start: + return + self.is_start = True + + async def _run_task_loop(): + while True: + logger.info("Dall E node is waiting for task...") + task = await self.task_queue.get() + logger.info(f"Dall E node get task: {task.display()}") + result = self._run_task(task) + + asyncio.create_task(_run_task_loop()) + + def display(self) -> str: + return f"DallE_ComputeNode: {self.node_id}" + + def get_task_state(self, task_id: str): + pass + + def get_capacity(self): + pass + + def is_support(self, task: ComputeTask) -> bool: + return task.task_type == ComputeTaskType.TEXT_2_IMAGE + + def is_local(self) -> bool: + return False diff --git a/src/aios/environment/workflow_env.py b/src/aios/environment/workflow_env.py index c5dcce6..c8f8448 100644 --- a/src/aios/environment/workflow_env.py +++ b/src/aios/environment/workflow_env.py @@ -81,15 +81,6 @@ class CalenderEnvironment(Environment): 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("get_contact", "get contact info", self._get_contact,{"name":"name of contact"})) @@ -317,12 +308,10 @@ class PaintEnvironment(Environment): self.is_run = False paint_param = { - "prompt": "Keywords of the content of the painting", - "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" + "prompt": "Description of the content of the painting", } self.add_ai_function(SimpleAIFunction("paint", - "Draw a picture according to the keywords", + "Draw a picture according to the description", self._paint,paint_param)) def _do_get_value(self,key:str) -> Optional[str]: diff --git a/test/test_dall_e_node.py b/test/test_dall_e_node.py new file mode 100644 index 0000000..37c1d62 --- /dev/null +++ b/test/test_dall_e_node.py @@ -0,0 +1,47 @@ +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.dall_e_compute_node import DallE_ComputeNode +from aios_kernel.compute_task import ComputeTaskType, ComputeTask, ComputeTaskState + +os.environ["TEXT2IMG_OUTPUT_DIR"] = "./" +os.environ["openai_api_key"] = "" + +@pytest.mark.asyncio +async def test_dall_e_node(propmt, model): + node = DallE_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['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) + +if __name__ == "__main__": + arg_len = len(os.sys.argv) + prompt = "a beautiful sunset" + model = "dall-e-3" + if arg_len >= 2: + prompt = os.sys.argv[1] + if arg_len == 3: + model = os.sys.argv[2] + + asyncio.run(test_dall_e_node(prompt, model)) diff --git a/test/test_local_sd_node.py b/test/test_local_sd_node.py index 21f8219..4404670 100644 --- a/test/test_local_sd_node.py +++ b/test/test_local_sd_node.py @@ -14,12 +14,12 @@ 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"] = "" -os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly" -os.environ["TEXT2IMG_OUTPUT_DIR"] = "./" +#os.environ["LOCAL_STABILITY_URL"] = "http://aigc:7860" +#os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly" +#os.environ["TEXT2IMG_OUTPUT_DIR"] = "./" @pytest.mark.asyncio -async def test_local_sd_node(propmt, model): +async def test_local_sd_node(propmt, model, negative_prompt): node = Local_Stability_ComputeNode.get_instance() if await node.initial() is not True: print("node initial failed!") @@ -31,6 +31,7 @@ async def test_local_sd_node(propmt, model): task.task_id = uuid.uuid4().hex task.params['model_name'] = model task.params['prompt'] = propmt + task.params['negative_prompt'] = negative_prompt await node.push_task(task) while True: @@ -51,9 +52,12 @@ if __name__ == "__main__": arg_len = len(os.sys.argv) prompt = "a beautiful sunset" model = "v1-5-pruned-emaonly" + negative_prompt = None if arg_len >= 2: prompt = os.sys.argv[1] if arg_len == 3: model = os.sys.argv[2] + if arg_len == 4: + negative_prompt = os.sys.argv[3] - asyncio.run(test_local_sd_node(prompt, model)) + asyncio.run(test_local_sd_node(prompt, model, negative_prompt))