Merge pull request #45 from glen0125/MVP

Add stability node (Text2Img)

src/aios_kernel/compute_kernel.py

@@ -13,8 +13,11 @@ logger = logging.getLogger(__name__)
 # How to dispatch different computing tasks (some tasks may contain a large amount of state for correct execution)
 # to suitable computing nodes, achieving a balance of speed, cost, and power consumption,
 # is the CORE GOAL of the entire computing task schedule system (aios_kernel).
+
+
 class ComputeKernel:
     _instance = None
+
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super().__new__(cls)
@@ -36,18 +39,19 @@ class ComputeKernel:
     def run(self, task: ComputeTask) -> None:
         # check there is compute node can support this task
         if self.is_task_support(task) is False:
-            logger.error(f"task {task.display()} is not support by any compute node")
+            logger.error(
+                f"task {task.display()} is not support by any compute node")
             return
         # add task to working_queue
         self.task_queue.put_nowait(task)
 
-
     def start(self):
         if self.is_start is True:
             logger.warn("compute_kernel is already start")
             return
 
         self.is_start = True
+
         async def _run_task_loop():
             while True:
                 logger.info("compute_kernel is waiting for task...")
@@ -60,17 +64,18 @@ class ComputeKernel:
 
         asyncio.create_task(_run_task_loop())
 
-
     def _schedule(self, task) -> ComputeNode:
         for node in self.compute_nodes.values():
-            if node.is_support(task) is True:
+            if node.is_support(task.task_type) is True:
                 return node
-        logger.warning(f"task {task.display()} is not support by any compute node")
+        logger.warning(
+            f"task {task.display()} is not support by any compute node")
         return None
 
     def add_compute_node(self, node: ComputeNode):
         if self.compute_nodes.get(node.node_id) is not None:
-            logger.warn(f"compute_node {node.display()} already in compute_kernel")
+            logger.warn(
+                f"compute_node {node.display()} already in compute_kernel")
             return
         self.compute_nodes[node.node_id] = node
         logger.info(f"add compute_node {node.display()} to compute_kernel")
@@ -85,7 +90,6 @@ class ComputeKernel:
     def is_task_support(self, task: ComputeTask) -> bool:
         return True
 
-
     # friendly interface for use:
     def llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0):
         # craete a llm_work_task ,push on queue's end
@@ -97,6 +101,7 @@ class ComputeKernel:
 
     async def do_llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0) -> str:
         task_req = self.llm_completion(prompt, mode_name, max_token)
+
         async def check_timer():
             check_times = 0
             while True:
@@ -118,4 +123,3 @@ class ComputeKernel:
             return task_req.result.result_str
 
         return "error!"
-    

src/aios_kernel/compute_node.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from .compute_task import ComputeTask
+from .compute_task import ComputeTask, ComputeTaskType
 
 
 class ComputeNode(ABC):
@@ -28,7 +28,7 @@ class ComputeNode(ABC):
         pass
 
     @abstractmethod
-    def is_support(self,task_type:str) -> bool:
+    def is_support(self, task_type: ComputeTaskType) -> bool:
         pass
 
     @abstractmethod
@@ -42,12 +42,9 @@ class ComputeNode(ABC):
         return "free"
 
 
-
 class LocalComputeNode(ComputeNode):
     def display(self) -> str:
         return super().display()
 
     def is_local(self) -> bool:
         return True
-    
-

src/aios_kernel/compute_task.py

@@ -3,6 +3,7 @@ from enum import Enum
 import uuid
 import time
 
+
 class ComputeTaskState(Enum):
     DONE = 0
     INIT = 1
@@ -11,9 +12,18 @@ class ComputeTaskState(Enum):
     PENDING = 4
 
 
+class ComputeTaskType(Enum):
+    NONE = -1
+    LLM_COMPLETION = 0
+    TEXT_2_IMAGE = 1
+    IMAGE_2_IMAGE = 2
+    VOICE_2_TEXT = 3
+    TEXT_2_VOICE = 4
+
+
 class ComputeTask:
     def __init__(self) -> None:
-        self.task_type = "llm_completion"
+        self.task_type = ComputeTaskType.NONE
         self.create_time = None
 
         self.task_id: str = None
@@ -27,7 +37,7 @@ class ComputeTask:
         self.error_str = None
 
     def set_llm_params(self, prompts, model_name, max_token_size, callchain_id=None):
-        self.task_type = "llm_completion"
+        self.task_type = ComputeTaskType.LLM_COMPLETION
         self.create_time = time.time()
         self.task_id = uuid.uuid4().hex
         self.callchain_id = callchain_id

src/aios_kernel/open_ai_node.py

@@ -5,13 +5,15 @@ import asyncio
 from asyncio import Queue
 import logging
 
-from .compute_task import ComputeTask,ComputeTaskResult,ComputeTaskState
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
 from .compute_node import ComputeNode
 
 logger = logging.getLogger(__name__)
 
+
 class OpenAI_ComputeNode(ComputeNode):
     _instance = None
+
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super(OpenAI_ComputeNode, cls).__new__(cls)
@@ -91,20 +93,11 @@ class OpenAI_ComputeNode(ComputeNode):
     def get_task_state(self, task_id: str):
         pass
 
-
     def get_capacity(self):
         pass
 
-
+    def is_support(self, task_type: ComputeTaskType) -> bool:
-    def is_support(self,task_type:str) -> bool:
+        return task_type == ComputeTaskType.LLM_COMPLETION
-        return True
-
 
     def is_local(self) -> bool:
         return False
-
-
-
-    
-    
-    

src/aios_kernel/stability_node.py

@@ -0,0 +1,136 @@
+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
+
+    def __new__(cls):
+        if cls._instanace is None:
+            cls._instanace = super(Stability_ComputeNode, cls).__new__(cls)
+            cls._instanace.is_start = False
+        return cls._instanace
+
+    def __init__(self) -> None:
+        super().__init__()
+        if self.is_start is True:
+            logger.warn("Stability_ComputeNode is already start")
+            return
+
+        self.is_start = True
+        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("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):
+        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

src/requirements.txt

@@ -0,0 +1,8 @@
+aiofiles==23.2.1
+aiohttp==3.8.5
+openai==0.28.0
+Pillow==10.0.0
+Pillow==10.0.0
+prompt_toolkit==3.0.39
+stability_sdk==0.8.4
+toml==0.10.2