Merge pull request #105 from glen0125/MVP

Add DallE node for text2imge

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. 
+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.
-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. 
+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.
-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. 
+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:
-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."""
+- 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."""

src/aios_kernel/__init__.py

@@ -27,6 +27,7 @@ from .compute_node_config import ComputeNodeConfig
 from .ai_function import SimpleAIFunction
 from .workspace_env import WorkspaceEnvironment
 from .openai_tts_node import OpenAITTSComputeNode   
+from .dall_e_compute_node import DallE_ComputeNode
 
 AIOS_Version = "0.5.2, build 2023-11-29"
 

src/aios_kernel/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

src/aios_kernel/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",
+            "prompt": "Description 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"
         }
         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]:

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))

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["LOCAL_STABILITY_URL"] = "http://aigc:7860"
-os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly"
+#os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly"
-os.environ["TEXT2IMG_OUTPUT_DIR"] = "./"
+#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))