Liu Zhicong
GitHub
@@ -0,0 +1,163 @@
|
||||
import os
|
||||
import io
|
||||
import asyncio
|
||||
from asyncio import Queue
|
||||
import logging
|
||||
import base64
|
||||
from PIL import Image
|
||||
import requests
|
||||
|
||||
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
|
||||
from .compute_node import ComputeNode
|
||||
from .storage import AIStorage, UserConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Local_Stability_ComputeNode(ComputeNode):
|
||||
_instance = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = Local_Stability_ComputeNode()
|
||||
return cls._instance
|
||||
|
||||
@classmethod
|
||||
def declare_user_config(cls):
|
||||
user_config = AIStorage.get_instance().get_user_config()
|
||||
user_config.add_user_config(
|
||||
"local_stability_url", "local stability url", False, None)
|
||||
user_config.add_user_config(
|
||||
"text2img_output_dir", "output dir", True, "./")
|
||||
user_config.add_user_config(
|
||||
"text2img_default_model", "text2img default model", True, "v1-5-pruned-emaonly")
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.is_start = False
|
||||
self.node_id = "local_stability_node"
|
||||
self.url = None
|
||||
self.default_model = None
|
||||
self.output_dir = None
|
||||
|
||||
self.task_queue = Queue()
|
||||
|
||||
async def initial(self):
|
||||
if os.getenv("LOCAL_STABILITY_URL") is not None:
|
||||
self.url = os.getenv("LOCAL_STABILITY_URL")
|
||||
else:
|
||||
self.url = AIStorage.get_instance(
|
||||
).get_user_config().get_value("local_stability_url")
|
||||
|
||||
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 os.getenv("TEXT2IMG_DEFAULT_MODEL") is not None:
|
||||
self.default_model = os.getenv("TEXT2IMG_DEFAULT_MODEL")
|
||||
else:
|
||||
self.default_model = AIStorage.get_instance(
|
||||
).get_user_config().get_value("text2img_default_model")
|
||||
|
||||
if self.url is None:
|
||||
logger.error("local stability url is None!")
|
||||
return False
|
||||
|
||||
if self.default_model is None:
|
||||
logger.error("local stability default model is None!")
|
||||
return False
|
||||
|
||||
if self.output_dir is None:
|
||||
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
|
||||
model_name = task.params["model_name"]
|
||||
prompts = task.params["prompts"]
|
||||
|
||||
logging.info(f"call local stability {model_name} prompts: {prompts}")
|
||||
|
||||
if model_name is not None:
|
||||
payload = {
|
||||
"sd_model_checkpoint": model_name,
|
||||
}
|
||||
# {'error': 'RuntimeError', 'detail': '', 'body': '', 'errors': "model 'xxx' not found"}
|
||||
response = requests.post(
|
||||
url=f'{self.url}/sdapi/v1/options', json=payload)
|
||||
if response.status_code != 200:
|
||||
task.state = ComputeTaskState.ERROR
|
||||
logger.error(
|
||||
f"set local stability model failed. err:{response.json()['errors']}")
|
||||
return None
|
||||
|
||||
payload = {
|
||||
"prompt": prompts,
|
||||
"steps": 20
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
url=f'{self.url}/sdapi/v1/txt2img', json=payload)
|
||||
r = response.json()
|
||||
|
||||
print(len(r['images']))
|
||||
for i in r['images']:
|
||||
image = Image.open(io.BytesIO(
|
||||
base64.b64decode(i.split(",", 1)[0])))
|
||||
file_name = os.path.join(self.output_dir, task.task_id + ".png")
|
||||
image.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("local_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
|
||||
@@ -17,6 +17,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
class Stability_ComputeNode(ComputeNode):
|
||||
_instanace = None
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls):
|
||||
if cls._instance is None:
|
||||
@@ -26,33 +27,47 @@ class Stability_ComputeNode(ComputeNode):
|
||||
@classmethod
|
||||
def declare_user_config(cls):
|
||||
user_config = AIStorage.get_instance().get_user_config()
|
||||
user_config.add_user_config("stability_api_key",False,None,"STABILITY_API_KEY")
|
||||
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")
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self.is_start = False
|
||||
self.node_id = "stability_node"
|
||||
self.api_key = ""
|
||||
self.engine = "stable-diffusion-512-v2-1"
|
||||
self.model = ""
|
||||
|
||||
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 = 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_ENGINE") is not None:
|
||||
self.engine = os.getenv("STABILITY_ENGINE")
|
||||
if os.getenv("STABILITY_MODEL") is not None:
|
||||
self.model = os.getenv("STABILITY_MODEL")
|
||||
else:
|
||||
self.model = AIStorage.get_instance().get_user_config().get_value("stability_model")
|
||||
|
||||
self.client = client.StabilityInference(
|
||||
key=self.api_key,
|
||||
verbose=True, # Print debug messages.
|
||||
engine=self.engine,
|
||||
engine=self.model,
|
||||
)
|
||||
|
||||
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)
|
||||
@@ -65,7 +80,7 @@ class Stability_ComputeNode(ComputeNode):
|
||||
# model_name && max_token_size not used here
|
||||
prompts = task.params["prompts"]
|
||||
|
||||
logging.info(f"call stability {self.engine} prompts: {prompts}")
|
||||
logging.info(f"call stability {self.model} prompts: {prompts}")
|
||||
answers = self.client.generate(
|
||||
prompt=prompts,
|
||||
# If a seed is provided, the resulting generated image will be deterministic.
|
||||
@@ -90,7 +105,8 @@ class Stability_ComputeNode(ComputeNode):
|
||||
|
||||
for resp in answers:
|
||||
for artifact in resp.artifacts:
|
||||
logger.info(f"artifact:{artifact.id},{artifact.type},{artifact.finish_reason}")
|
||||
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")
|
||||
@@ -113,6 +129,7 @@ class Stability_ComputeNode(ComputeNode):
|
||||
if self.is_start:
|
||||
return
|
||||
self.is_start = True
|
||||
|
||||
async def _run_task_loop():
|
||||
while True:
|
||||
logger.info("stability_node is waiting for task...")
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
import io
|
||||
import asyncio
|
||||
import sys
|
||||
import logging
|
||||
import pytest
|
||||
directory = os.path.dirname(__file__)
|
||||
sys.path.append(directory + '/../src')
|
||||
from aios_kernel.local_stability_node import Local_Stability_ComputeNode
|
||||
from aios_kernel.compute_task import ComputeTaskType, ComputeTask, ComputeTaskState
|
||||
|
||||
#to launch a local stability node, please check:
|
||||
#https://github.com/glen0125/stable-diffusion-webui-docker
|
||||
|
||||
os.environ["LOCAL_STABILITY_URL"] = "http://aigc:7866"
|
||||
os.environ["TEXT2IMG_DEFAULT_MODEL"] = "v1-5-pruned-emaonly"
|
||||
os.environ["TEXT2IMG_OUTPUT_DIR"] = "./"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_sd_node(propmt, model):
|
||||
node = Local_Stability_ComputeNode.get_instance()
|
||||
if await node.initial() is not True:
|
||||
print("node initial failed!")
|
||||
return
|
||||
|
||||
task = ComputeTask()
|
||||
task.task_type = ComputeTaskType.TEXT_2_IMAGE
|
||||
task.create_time = time.time()
|
||||
task.task_id = uuid.uuid4().hex
|
||||
task.params['model_name'] = model
|
||||
task.params['prompts'] = propmt
|
||||
await node.push_task(task)
|
||||
|
||||
while True:
|
||||
if task.state == ComputeTaskState.DONE:
|
||||
local_file = task.result.result
|
||||
print("local file is: ", local_file)
|
||||
break
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# result = node._run_task(task)
|
||||
# print("result is: ", result)
|
||||
|
||||
# if result.result is not None:
|
||||
# local_file = result.result
|
||||
# print("local file is: ", local_file)
|
||||
|
||||
if __name__ == "__main__":
|
||||
arg_len = len(os.sys.argv)
|
||||
prompt = "a beautiful sunset"
|
||||
model = "v1-5-pruned-emaonly"
|
||||
if arg_len >= 2:
|
||||
prompt = os.sys.argv[1]
|
||||
if arg_len == 3:
|
||||
model = os.sys.argv[2]
|
||||
|
||||
asyncio.run(test_local_sd_node(prompt, model))
|
||||
Reference in New Issue
Block a user