opendan/src/aios_kernel/stability_node.py

import os
import io
import asyncio
from asyncio import Queue
import logging

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
from .compute_node import ComputeNode

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


    def __init__(self) -> None:
        super().__init__()

        self.is_start = False
        self.node_id = "stability_node"
        self.api_key = ""
        self.engine = "stable-diffusion-512-v2-1"

        self.task_queue = Queue()

        if os.getenv("STABILITY_API_KEY") is not None:
            self.api_key = os.getenv("STABILITY_API_KEY")

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

        self.client = client.StabilityInference(
            key=self.api_key,
            verbose=True,  # Print debug messages.
            engine=self.engine,
        )

        self.start()

    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 && max_token_size not used here
        prompts = task.params["prompts"]

        logging.info(f"call stability {self.engine} prompts: {prompts}")
        answers = self.client.generate(
            prompt=prompts,
            # 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:
                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:
                    img = Image.open(io.BytesIO(artifact.binary))
                    # Save our generated images with the task_id as the filename.
                    file_name = task.task_id + ".png"  # which dir to save?
                    img.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("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_type: ComputeTaskType) -> bool:
        return task_type == ComputeTaskType.TEXT_2_IMAGE

    def is_local(self) -> bool:
        return False