gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add stability node (Text2Img)

Browse Source
This commit is contained in:
Song
2023-09-07 12:50:13 +08:00
parent 1bbdf2e722
commit 3eac598c97
5 changed files with 209 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/compute_node.py
+11 -14
View File
@@ -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
src/aios_kernel/compute_task.py
+26 -16
View File
@@ -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
src/aios_kernel/open_ai_node.py
+24 -31
View File
@@ -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
src/aios_kernel/stability_node.py
+140
View File
@@ -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
src/requirements.txt
+8
View File
@@ -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