Add a local sd node

src/aios_kernel/local_stability_node.py

@@ -0,0 +1,163 @@
+import os
+import io
+import asyncio
+from asyncio import Queue
+import logging
+import base64
+from PIL import Image
+import requests
+
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
+from .compute_node import ComputeNode
+from .storage import AIStorage, UserConfig
+
+logger = logging.getLogger(__name__)
+
+
+class Local_Stability_ComputeNode(ComputeNode):
+    _instance = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = Local_Stability_ComputeNode()
+        return cls._instance
+
+    @classmethod
+    def declare_user_config(cls):
+        user_config = AIStorage.get_instance().get_user_config()
+        user_config.add_user_config(
+            "local_stability_url", "local stability url", False, None)
+        user_config.add_user_config(
+            "text2img_output_dir", "output dir", True, "./")
+        user_config.add_user_config(
+            "text2img_default_model", "text2img default model", True, "v1-5-pruned-emaonly")
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.is_start = False
+        self.node_id = "local_stability_node"
+        self.url = None
+        self.default_model = None
+        self.output_dir = None
+
+        self.task_queue = Queue()
+    
+    async def initial(self):
+        if os.getenv("LOCAL_STABILITY_URL") is not None:
+            self.url = os.getenv("LOCAL_STABILITY_URL")
+        else:
+            self.url = AIStorage.get_instance(
+            ).get_user_config().get_value("local_stability_url")
+
+        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 os.getenv("TEXT2IMG_DEFAULT_MODEL") is not None:
+            self.default_model = os.getenv("TEXT2IMG_DEFAULT_MODEL")
+        else:
+            self.default_model = AIStorage.get_instance(
+            ).get_user_config().get_value("text2img_default_model")
+
+        if self.url is None:
+            logger.error("local stability url is None!")
+            return False
+
+        if self.default_model is None:
+            logger.error("local stability default model is None!")
+            return False
+
+        if self.output_dir is None:
+            self.output_dir = "./"
+
+        self.start()
+
+        return True
+
+    async def push_task(self, task: ComputeTask, proiority: int = 0):
+        logger.info(f"stability_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
+        model_name = task.params["model_name"]
+        prompts = task.params["prompts"]
+
+        logging.info(f"call local stability {model_name} prompts: {prompts}")
+
+        if model_name is not None:
+            payload = {
+                "sd_model_checkpoint": model_name,
+            }
+            # {'error': 'RuntimeError', 'detail': '', 'body': '', 'errors': "model 'xxx' not found"}
+            response = requests.post(
+                url=f'{self.url}/sdapi/v1/options', json=payload)
+            if response.status_code != 200:
+                task.state = ComputeTaskState.ERROR
+                logger.error(
+                    f"set local stability model failed. err:{response.json()['errors']}")
+                return None
+
+        payload = {
+            "prompt": prompts,
+            "steps": 20
+        }
+
+        response = requests.post(
+            url=f'{self.url}/sdapi/v1/txt2img', json=payload)
+        r = response.json()
+
+        print(len(r['images']))
+        for i in r['images']:
+            image = Image.open(io.BytesIO(
+                base64.b64decode(i.split(",", 1)[0])))
+            file_name = os.path.join(self.output_dir, task.task_id + ".png")
+            image.save(file_name)
+
+            result = ComputeTaskResult()
+            result.set_from_task(task)
+            result.worker_id = self.node_id
+            result.result = {"file": file_name}
+
+            return result
+
+        return None
+
+    def start(self):
+        if self.is_start:
+            return
+        self.is_start = True
+
+        async def _run_task_loop():
+            while True:
+                logger.info("local_stability_node is waiting for task...")
+                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:
+                    task.state = ComputeTaskState.DONE
+                    task.result = result
+
+        asyncio.create_task(_run_task_loop())
+
+    def display(self) -> str:
+        return f"Stability_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/stability_node.py

@@ -10,49 +10,64 @@ import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
 
 from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
 from .compute_node import ComputeNode
-from .storage import AIStorage,UserConfig
+from .storage import AIStorage, UserConfig
 
 logger = logging.getLogger(__name__)
 
 
 class Stability_ComputeNode(ComputeNode):
     _instanace = None
+
     @classmethod
     def get_instance(cls):
         if cls._instance is None:
             cls._instance = Stability_ComputeNode()
         return cls._instance
-    
+
     @classmethod
     def declare_user_config(cls):
         user_config = AIStorage.get_instance().get_user_config()
-        user_config.add_user_config("stability_api_key",False,None,"STABILITY_API_KEY")
+        user_config.add_user_config(
+            "stability_api_key", "stability api key", False, None)
+        user_config.add_user_config(
+            "stability_model", "stability model name", True, "stable-diffusion-512-v2-1")
 
-    def __init__(self) -> None:
+    def __init__(self):
         super().__init__()
 
         self.is_start = False
         self.node_id = "stability_node"
         self.api_key = ""
-        self.engine = "stable-diffusion-512-v2-1"
+        self.model = ""
 
         self.task_queue = Queue()
 
         if os.getenv("STABILITY_API_KEY") is not None:
             self.api_key = os.getenv("STABILITY_API_KEY")
+        else:
+            self.api_key = AIStorage.get_instance(
+            ).get_user_config().get_value("stability_api_key")
+
+        if self.api_key is None:
+            logger.error("stability api key is None!")
+            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_ENGINE") is not None:
+        if os.getenv("STABILITY_MODEL") is not None:
-            self.engine = os.getenv("STABILITY_ENGINE")
+            self.model = os.getenv("STABILITY_MODEL")
+        else:
+            self.model = AIStorage.get_instance().get_user_config().get_value("stability_model")
 
         self.client = client.StabilityInference(
             key=self.api_key,
             verbose=True,  # Print debug messages.
-            engine=self.engine,
+            engine=self.model,
         )
 
         self.start()
 
+        return True
+
     async def push_task(self, task: ComputeTask, proiority: int = 0):
         logger.info(f"stability_node push task: {task.display()}")
         self.task_queue.put_nowait(task)
@@ -65,7 +80,7 @@ class Stability_ComputeNode(ComputeNode):
         # model_name && max_token_size not used here
         prompts = task.params["prompts"]
 
-        logging.info(f"call stability {self.engine} prompts: {prompts}")
+        logging.info(f"call stability {self.model} prompts: {prompts}")
         answers = self.client.generate(
             prompt=prompts,
             # If a seed is provided, the resulting generated image will be deterministic.
@@ -90,8 +105,9 @@ class Stability_ComputeNode(ComputeNode):
 
         for resp in answers:
             for artifact in resp.artifacts:
-                logger.info(f"artifact:{artifact.id},{artifact.type},{artifact.finish_reason}")
+                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")
                 if artifact.type == generation.ARTIFACT_IMAGE:
@@ -113,6 +129,7 @@ class Stability_ComputeNode(ComputeNode):
         if self.is_start:
             return
         self.is_start = True
+
         async def _run_task_loop():
             while True:
                 logger.info("stability_node is waiting for task...")

test/test_local_sd_node.py

@@ -0,0 +1,59 @@
+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.local_stability_node import Local_Stability_ComputeNode
+from aios_kernel.compute_task import ComputeTaskType, ComputeTask, ComputeTaskState
+
+#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["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly"
+os.environ["TEXT2IMG_OUTPUT_DIR"] = "./"
+
+@pytest.mark.asyncio
+async def test_local_sd_node(propmt, model):
+    node = Local_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['prompts'] = 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 = "v1-5-pruned-emaonly"
+    if arg_len >= 2:
+        prompt = os.sys.argv[1]
+    if arg_len == 3:
+        model = os.sys.argv[2]
+
+    asyncio.run(test_local_sd_node(prompt, model))