Merge pull request #67 from glen0125/MVP

Return result code from sd node
This commit is contained in:
Liu Zhicong
2023-09-27 11:48:11 -07:00
committed by GitHub
6 changed files with 70 additions and 46 deletions
+10 -8
View File
@@ -175,17 +175,19 @@ class ComputeKernel:
return task_result.result
def text_2_image(self, prompt:str, model_name:Optional[str] = None):
def text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None):
task = ComputeTask()
task.set_text_2_image_params(prompt,model_name)
task.set_text_2_image_params(prompt,model_name, negative_prompt)
self.run(task)
return task
async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None) -> [str, ComputeTaskResult]:
task_req = self.text_2_image(prompt,model_name)
task_result = self._send_task(task_req)
if task_req.state == ComputeTaskState.DONE:
return None, task_result
async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult:
task = self.text_2_image(prompt,model_name, negative_prompt)
task = await self._send_task(task)
return task_req.error_str, None
return task.result
# if task_req.state == ComputeTaskState.DONE:
# return None, task_result
# return task_req.error_str, None
+2 -1
View File
@@ -71,12 +71,13 @@ class ComputeTask:
self.params["model_name"] = "text-embedding-ada-002"
self.params["input"] = input
def set_text_2_image_params(self, prompt: str, model_name, callchain_id=None):
def set_text_2_image_params(self, prompt: str, model_name, negative_prompt="", callchain_id=None):
self.task_type = ComputeTaskType.TEXT_2_IMAGE
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
self.params["prompt"] = prompt
self.params["negative_prompt"] = negative_prompt
if model_name is not None:
self.params["model_name"] = model_name
else:
+23 -10
View File
@@ -9,7 +9,7 @@ import requests
from typing import Tuple
from pathlib import Path
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType, ComputeTaskResultCode
from .compute_node import ComputeNode
from .storage import AIStorage, UserConfig
@@ -107,10 +107,19 @@ class Local_Stability_ComputeNode(ComputeNode):
def _run_task(self, task: ComputeTask):
task.state = ComputeTaskState.RUNNING
result = ComputeTaskResult()
result.result_code = ComputeTaskResultCode.ERROR
result.set_from_task(task)
model_name = task.params["model_name"]
prompt = task.params["prompt"]
negative_prompt = task.params["negative_prompt"]
if negative_prompt == None or negative_prompt == "":
negative_prompt = "sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, duplicate, mutated hands, mutated legs, (blurry:1.3), (bad anatomy:1.2), bad proportions, extra limbs, more than 2 nipples, extra legs, fused fingers, missing fingers, jpeg artifacts, signature, watermark, username, artist name, heterochromia, muscular legs, monochrome, grayscale, skin spots, acnes, skin blemishes, age spot, skin spots, acnes, logo, badhandv4, easynegative, cropped image, patreon,lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, ng_deepnegative_v1_75t, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry,(Tiptoe:1.3),looking at viewer, Twisted eyes"
logging.info(f"call local stability {model_name} prompts: {prompt}")
prompt += ",masterpiece, best quality:1.3"
logging.info(f"call local stability {model_name} prompts: {prompt}, nagative_prompt: {negative_prompt}")
if model_name is not None:
payload = {
@@ -123,12 +132,14 @@ class Local_Stability_ComputeNode(ComputeNode):
err_msg = f"Set local stability model failed. err:{err}"
logger.error(err_msg)
task.error_str = err_msg
return None
result.error_str = err_msg
return result
logging.info(f"set local stability model {model_name} success")
payload = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"steps": 20
}
@@ -138,7 +149,8 @@ class Local_Stability_ComputeNode(ComputeNode):
err_msg = f"Failed. err:{err}"
logger.error(err_msg)
task.error_str = err_msg
return None
result.error_str = err_msg
return result
r = resp.json()
@@ -148,16 +160,17 @@ class Local_Stability_ComputeNode(ComputeNode):
file_name = os.path.join(self.output_dir, task.task_id + ".png")
image.save(file_name)
result = ComputeTaskResult()
result.set_from_task(task)
task.state = ComputeTaskState.DONE
result.result_code = ComputeTaskResultCode.OK
result.worker_id = self.node_id
result.result = {"file": file_name}
return result
task.error_str = "Unknown error!"
result.error_str = "Unknown error!"
task.state = ComputeTaskState.ERROR
return None
return result
def start(self):
if self.is_start:
@@ -170,9 +183,9 @@ class Local_Stability_ComputeNode(ComputeNode):
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
# if result is not None:
# task.state = ComputeTaskState.DONE
# task.result = result
asyncio.create_task(_run_task_loop())
+18 -10
View File
@@ -9,7 +9,7 @@ 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_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType, ComputeTaskResultCode
from .compute_node import ComputeNode
from .storage import AIStorage, UserConfig
@@ -95,10 +95,15 @@ class Stability_ComputeNode(ComputeNode):
def _run_task(self, task: ComputeTask):
task.state = ComputeTaskState.RUNNING
result = ComputeTaskResult()
result.result_code = ComputeTaskResultCode.ERROR
result.set_from_task(task)
model_name = task.params["model_name"]
prompt = task.params["prompt"]
negative_prompt = task.params["negative_prompt"]
logging.info(f"call stability {self.default_model} prompts: {prompt}")
logging.info(f"call stability {self.default_model} prompts: {prompt}, negative_prompt: {negative_prompt}")
api = None
try:
@@ -109,9 +114,10 @@ class Stability_ComputeNode(ComputeNode):
)
except Exception as e:
task.error_str = f"create stability client failed: {e}"
result.error_str = f"create stability client failed: {e}"
logging.warn(task.error_str)
task.state = ComputeTaskState.ERROR
return None
return result
answers = api.generate(
prompt=prompt,
@@ -141,24 +147,26 @@ class Stability_ComputeNode(ComputeNode):
err_msg = "request activated the API's safety filters"
logging.warn(err_msg)
task.error_str = err_msg
result.error_str = err_msg
task.state = ComputeTaskState.ERROR
return None
return result
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 = os.path.join(self.output_dir, task.task_id + ".png")
img.save(file_name)
result = ComputeTaskResult()
result.set_from_task(task)
task.state = ComputeTaskState.DONE
result.result_code = ComputeTaskResultCode.OK
result.worker_id = self.node_id
result.result = {"file": file_name}
return result
task.error_str = "Unknown error!"
result.error_str = "Unknown error!"
task.state = ComputeTaskState.ERROR
return None
return result
def start(self):
if self.is_start:
@@ -171,9 +179,9 @@ class Stability_ComputeNode(ComputeNode):
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
# if result is not None:
# task.state = ComputeTaskState.DONE
# task.result = result
asyncio.create_task(_run_task_loop())
+10 -9
View File
@@ -9,7 +9,7 @@ import logging
from typing import Optional
from .text_to_speech_function import TextToSpeechFunction
from .compute_kernel import ComputeKernel
from .compute_kernel import ComputeKernel, ComputeTaskResultCode
from .environment import Environment,EnvironmentEvent
from .ai_function import SimpleAIFunction
from .storage import AIStorage
@@ -311,9 +311,9 @@ class CalenderEnvironment(Environment):
async def _paint(self, prompt, model_name = None) -> str:
err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name)
if err is not None:
return f"exec paint failed. err:{err}"
result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name)
if result.result_code == ComputeTaskResultCode.ERROR:
return f"exec paint failed. err:{result.error_str}"
else:
return f'exec paint OK, saved as a local file, path is: {result.result["file"]}'
@@ -324,19 +324,20 @@ class PaintEnvironment(Environment):
self.is_run = False
paint_param = {
"prompt": "A description of the content of the painting",
"model_name": "Which model to use to draw the picture, can be None"
"prompt": "Keywords of the content of the painting",
"model_name": "Which model to use to draw the picture, can be None",
"negative_prompt": "Keywords that describe what is not to be drawn, can be None"
}
self.add_ai_function(SimpleAIFunction("paint",
"Draw a picture according to the description",
"Draw a picture according to the keywords",
self._paint,paint_param))
def _do_get_value(self,key:str) -> Optional[str]:
return None
async def _paint(self, prompt, model_name = None) -> str:
err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name)
async def _paint(self, prompt, model_name = None, negative_prompt = None) -> str:
err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name, negative_prompt)
if err is not None:
return f"exec paint failed. err:{err}"
else: