import os import io import asyncio from asyncio import Queue import logging import base64 from PIL import Image import requests from typing import Tuple from pathlib import Path from aios import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType,ComputeTaskResultCode,ComputeNode,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() if os.getenv("LOCAL_STABILITY_URL") is None: user_config.add_user_config( "local_stability_url", "local stability url", True, None) 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("TEXT2IMG_DEFAULT_MODEL") is None: 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.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"stability_node push task: {task.display()}") self.task_queue.put_nowait(task) async def remove_task(self, task_id: str): pass def _make_post_request(self, url, json) -> Tuple[str, requests.Response]: try: response = requests.post(url, json=json) if response.status_code != 200: return f'{response.status_code}, {response.json()}', None return None, response except Exception as e: return f"{e}", None def _run_task(self, task: ComputeTask): task.state = ComputeTaskState.RUNNING result = ComputeTaskResult() result.result_code = ComputeTaskResultCode.ERROR result.set_from_task(task) model_name = task.params["model_name"] prompt = task.params["prompt"] negative_prompt = task.params["negative_prompt"] if negative_prompt == None or negative_prompt == "": negative_prompt = "sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, duplicate, mutated hands, mutated legs, (blurry:1.3), (bad anatomy:1.2), bad proportions, extra limbs, more than 2 nipples, extra legs, fused fingers, missing fingers, jpeg artifacts, signature, watermark, username, artist name, heterochromia, muscular legs, monochrome, grayscale, skin spots, acnes, skin blemishes, age spot, skin spots, acnes, logo, badhandv4, easynegative, cropped image, patreon,lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, ng_deepnegative_v1_75t, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry,(Tiptoe:1.3),looking at viewer, Twisted eyes" prompt += ",masterpiece, best quality:1.3" logging.info(f"call local stability {model_name} prompts: {prompt}, nagative_prompt: {negative_prompt}") if model_name is not None: payload = { "sd_model_checkpoint": model_name, } err, resp = self._make_post_request(f'{self.url}/sdapi/v1/options', payload) if err is not None: task.state = ComputeTaskState.ERROR err_msg = f"Set local stability model failed. err:{err}" logger.error(err_msg) task.error_str = err_msg result.error_str = err_msg return result logging.info(f"set local stability model {model_name} success") payload = { "prompt": prompt, "negative_prompt": negative_prompt, "steps": 20 } err, resp = self._make_post_request(f'{self.url}/sdapi/v1/txt2img', payload) if err is not None: task.state = ComputeTaskState.ERROR err_msg = f"Failed. err:{err}" logger.error(err_msg) task.error_str = err_msg result.error_str = err_msg return result r = resp.json() 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) task.state = ComputeTaskState.DONE result.result_code = ComputeTaskResultCode.OK result.worker_id = self.node_id result.result = {"file": file_name} return result task.error_str = "Unknown error!" result.error_str = "Unknown error!" task.state = ComputeTaskState.ERROR return result 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