From 1971a9399d00995da542b8908e9bf220d2750cb8 Mon Sep 17 00:00:00 2001 From: Song Date: Wed, 20 Sep 2023 22:31:33 +0800 Subject: [PATCH] Add a local sd node --- src/aios_kernel/local_stability_node.py | 163 ++++++++++++++++++++++++ src/aios_kernel/stability_node.py | 39 ++++-- test/test_local_sd_node.py | 59 +++++++++ 3 files changed, 250 insertions(+), 11 deletions(-) create mode 100644 src/aios_kernel/local_stability_node.py create mode 100644 test/test_local_sd_node.py diff --git a/src/aios_kernel/local_stability_node.py b/src/aios_kernel/local_stability_node.py new file mode 100644 index 0000000..992e00f --- /dev/null +++ b/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 diff --git a/src/aios_kernel/stability_node.py b/src/aios_kernel/stability_node.py index 523ca17..b36e86c 100644 --- a/src/aios_kernel/stability_node.py +++ b/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: - self.engine = os.getenv("STABILITY_ENGINE") + if os.getenv("STABILITY_MODEL") is not None: + 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...") diff --git a/test/test_local_sd_node.py b/test/test_local_sd_node.py new file mode 100644 index 0000000..62659d1 --- /dev/null +++ b/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))