import os import io import asyncio from asyncio import Queue import logging from pathlib import Path from PIL import Image from stability_sdk import client import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType, ComputeTaskResultCode from .compute_node import ComputeNode from .storage import AIStorage, UserConfig logger = logging.getLogger(__name__) class Stability_ComputeNode(ComputeNode): _instance = 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", "stability api key", False, None) user_config.add_user_config( "stability_model", "stability model name", True, "stable-diffusion-512-v2-1") 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("STABILITY_DEFAULT_MODEL") is None: user_config.add_user_config( "stability_default_model", "stability default model", True, "stable-diffusion-512-v2-1") def __init__(self): super().__init__() self.is_start = False self.node_id = "stability_node" self.api_key = "" self.default_model = "" self.task_queue = Queue() async def initial(self): 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_DEFAULT_MODEL") is not None: self.default_model = os.getenv("STABILITY_DEFAULT_MODEL") else: self.default_model = AIStorage.get_instance().get_user_config().get_value("stability_default_model") if self.default_model is None: self.default_model = "stable-diffusion-512-v2-1" 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"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 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"] logging.info(f"call stability {self.default_model} prompts: {prompt}, negative_prompt: {negative_prompt}") api = None try: api = client.StabilityInference( key=self.api_key, verbose=True, # Print debug messages. engine=model_name, ) except Exception as e: task.error_str = f"create stability client failed: {e}" result.error_str = f"create stability client failed: {e}" logging.warn(task.error_str) task.state = ComputeTaskState.ERROR return result answers = api.generate( prompt=prompt, # If a seed is provided, the resulting generated image will be deterministic. seed=0, # What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again. # Note: This isn't quite the case for Clip Guided generations, which we'll tackle in a future example notebook. # Amount of inference steps performed on image generation. Defaults to 30. steps=30, # Influences how strongly your generation is guided to match your prompt. cfg_scale=7.0, # Setting this value higher increases the strength in which it tries to match your prompt. # Defaults to 7.0 if not specified. width=512, # Generation width, defaults to 512 if not included. height=512, # Generation height, defaults to 512 if not included. # Number of images to generate, defaults to 1 if not included. samples=1, # Choose which sampler we want to denoise our generation with. sampler=generation.SAMPLER_K_DPMPP_2M # Defaults to k_dpmpp_2m if not specified. Clip Guidance only supports ancestral samplers. # (Available Samplers: ddim, plms, k_euler, k_euler_ancestral, k_heun, k_dpm_2, k_dpm_2_ancestral, k_dpmpp_2s_ancestral, k_lms, k_dpmpp_2m, k_dpmpp_sde) ) for resp in answers: for artifact in resp.artifacts: if artifact.finish_reason == generation.FILTER: err_msg = "request activated the API's safety filters" logging.warn(err_msg) task.error_str = err_msg result.error_str = err_msg task.state = ComputeTaskState.ERROR return result if artifact.type == generation.ARTIFACT_IMAGE: img = Image.open(io.BytesIO(artifact.binary)) # Save our generated images with the task_id as the filename. file_name = os.path.join(self.output_dir, task.task_id + ".png") img.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("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