diff --git a/src/aios_kernel/compute_node.py b/src/aios_kernel/compute_node.py index 6fe0059..7e159f7 100644 --- a/src/aios_kernel/compute_node.py +++ b/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): @@ -8,15 +8,15 @@ class ComputeNode(ABC): self.enable = True @abstractmethod - async def push_task(self,task:ComputeTask,proiority:int = 0): - pass - - @abstractmethod - async def remove_task(self,task_id:str): + async def push_task(self, task: ComputeTask, proiority: int = 0): pass @abstractmethod - def get_task_state(self,task_id:str): + async def remove_task(self, task_id: str): + pass + + @abstractmethod + def get_task_state(self, task_id: str): pass @abstractmethod @@ -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 @@ -37,17 +37,14 @@ class ComputeNode(ABC): def is_trusted(self) -> bool: return True - + def get_fee_type(self) -> str: return "free" - - + class LocalComputeNode(ComputeNode): def display(self) -> str: return super().display() - + def is_local(self) -> bool: return True - - diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index f2253e0..7895283 100644 --- a/src/aios_kernel/compute_task.py +++ b/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,22 +12,31 @@ 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.create_time = None - self.task_id:str = None - self.callchain_id:str = None - self.params:dict = {} - self.refers:dict = None - self.pading_data:bytearray = None + self.task_id: str = None + self.callchain_id: str = None + self.params: dict = {} + self.refers: dict = None + self.pading_data: bytearray = None self.state = ComputeTaskState.INIT self.result = None self.error_str = None - def set_llm_params(self,prompts,model_name,max_token_size,callchain_id = None): + def set_llm_params(self, prompts, model_name, max_token_size, callchain_id=None): self.task_type = "llm_completion" self.create_time = time.time() self.task_id = uuid.uuid4().hex @@ -34,7 +44,7 @@ class ComputeTask: self.params["prompts"] = prompts.messages if model_name is not None: self.params["model_name"] = model_name - else: + else: self.params["model_name"] = "gpt-4-0613" self.params["max_token_size"] = max_token_size @@ -45,16 +55,16 @@ class ComputeTask: class ComputeTaskResult: def __init__(self) -> None: self.create_time = None - self.task_id:str = None - self.callchain_id:str = None - self.worker_id:str = None - self.result_code:int = 0 - self.result_str:str = None + self.task_id: str = None + self.callchain_id: str = None + self.worker_id: str = None + self.result_code: int = 0 + self.result_str: str = None - self.result:dict = {} - self.result_refers:dict = None - self.pading_data:bytearray = None + self.result: dict = {} + self.result_refers: dict = None + self.pading_data: bytearray = None - def set_from_task(self,task:ComputeTask): + def set_from_task(self, task: ComputeTask): self.task_id = task.task_id self.callchain_id = task.callchain_id diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 7f5c8e6..06e0f09 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -5,47 +5,49 @@ 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) cls._instance.is_start = False return cls._instance - + def __init__(self) -> None: super().__init__() if self.is_start is True: logger.warn("OpenAI_ComputeNode is already start") return - + self.is_start = True - #openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud + # openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud self.openai_api_key = "" self.node_id = "openai_node" self.task_queue = Queue() - if os.getenv("OPENAI_API_KEY") is not None: + if os.getenv("OPENAI_API_KEY") is not None: openai.api_key = os.getenv("OPENAI_API_KEY") else: openai.api_key = self.openai_api_key - + self.start() - - async def push_task(self,task:ComputeTask,proiority:int = 0): + + async def push_task(self, task: ComputeTask, proiority: int = 0): logger.info(f"openai_node push task: {task.display()}") self.task_queue.put_nowait(task) - - async def remove_task(self,task_id:str): + + async def remove_task(self, task_id: str): pass - - def _run_task(self,task:ComputeTask): + + def _run_task(self, task: ComputeTask): task.state = ComputeTaskState.RUNNING mode_name = task.params["model_name"] # max_token_size = task.params["max_token_size"] @@ -57,19 +59,19 @@ class OpenAI_ComputeNode(ComputeNode): max_tokens=4000, temperature=1.2) logger.info(f"openai response: {resp}") - + status_code = resp["choices"][0]["finish_reason"] if status_code != "stop": task.state = ComputeTaskState.ERROR - task.error_str =f"The status code was {status_code}." + task.error_str = f"The status code was {status_code}." return None - - result = ComputeTaskResult() + + result = ComputeTaskResult() result.set_from_task(task) result.worker_id = self.node_id result.result_str = resp["choices"][0]["message"]["content"] result.result = resp["choices"][0]["message"] - + return result def start(self): @@ -82,29 +84,20 @@ class OpenAI_ComputeNode(ComputeNode): if result is not None: task.state = ComputeTaskState.DONE task.result = result - + asyncio.create_task(_run_task_loop()) def display(self) -> str: return f"OpenAI_ComputeNode: {self.node_id}" - - def get_task_state(self,task_id:str): - pass + def get_task_state(self, task_id: str): + pass def get_capacity(self): pass - - def is_support(self,task_type:str) -> bool: - return True - + def is_support(self, task_type: ComputeTaskType) -> bool: + return task_type == ComputeTaskType.LLM_COMPLETION def is_local(self) -> bool: return False - - - - - - \ No newline at end of file diff --git a/src/aios_kernel/stability_node.py b/src/aios_kernel/stability_node.py new file mode 100644 index 0000000..a5ed4e0 --- /dev/null +++ b/src/aios_kernel/stability_node.py @@ -0,0 +1,140 @@ +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 = "" # "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0" + self.engine = "" # stable-diffusion-512-v2-0 + + 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 = "sk-RQDlJtBFQg6I3IueeGCGZTPhWPYAl3bgRdvFDMkcEXsKbUc0" + + # 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") + else: + self.engine = "stable-diffusion-512-v2-1" + + 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 diff --git a/src/requirements.txt b/src/requirements.txt new file mode 100644 index 0000000..605dc72 --- /dev/null +++ b/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