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

Merge pull request #63 from glen0125/MVP

Browse Source 
Add a paint agent(David)
This commit is contained in:
Liu Zhicong
2023-09-23 01:42:13 -07:00
committed by GitHub
parent 2662f8569f 2be2ef3f60
commit 58572827b8
10 changed files with 136 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/aios_kernel/__init__.py
+1
View File
@@ -21,6 +21,7 @@ from .storage import ResourceLocation,AIStorage,UserConfig,UserConfigItem
from .contact_manager import ContactManager,Contact,FamilyMember
from .text_to_speech_function import TextToSpeechFunction
from .workspace_env import WorkspaceEnvironment
from .local_stability_node import Local_Stability_ComputeNode
AIOS_Version = "0.5.1, build 2023-9-17"
src/aios_kernel/agent.py
+2 -2
View File
@@ -123,7 +123,7 @@ class AIAgent:
self.unread_msg = Queue() # msg from other agent
self.owner_env : Environment = None
self.owenr_bus = None
self.enable_function_list = []
self.enable_function_list = None
@classmethod
def create_from_templete(cls,templete:AIAgentTemplete, fullname:str):
@@ -392,7 +392,7 @@ class AIAgent:
prompt.append(msg_prompt)
logger.debug(f"Agent {self.agent_id} do llm token static system:{system_prompt_len},function:{function_token_len},history:{history_token_len},input:{input_len}，totoal prompt:{system_prompt_len + function_token_len + history_token_len} ")
logger.debug(f"Agent {self.agent_id} do llm token static system:{system_prompt_len},function:{function_token_len},history:{history_token_len},input:{input_len}, totoal prompt:{system_prompt_len + function_token_len + history_token_len} ")
task_result:ComputeTaskResult = await ComputeKernel.get_instance().do_llm_completion(prompt,self.llm_model_name,self.max_token_size,inner_functions)
final_result = task_result.result_str
src/aios_kernel/compute_kernel.py
+32
View File
@@ -197,3 +197,35 @@ class ComputeKernel:
raise Exception("do_text_to_speech failed!")
def text_2_image(self, prompt:str, model_name:Optional[str] = None):
task = ComputeTask()
task.set_text_2_image_params(prompt,model_name)
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)
async def check_timer():
check_times = 0
while True:
if task_req.state == ComputeTaskState.DONE:
break
if task_req.state == ComputeTaskState.ERROR:
break
#may be it will take a long time to generate a image
if check_times >= 60:
task_req.state = ComputeTaskState.ERROR
break
await asyncio.sleep(0.5)
check_times += 1
await asyncio.create_task(check_timer())
if task_req.state == ComputeTaskState.DONE:
return None, task_req.result
return task_req.error_str, None
src/aios_kernel/compute_task.py
+11
View File
@@ -64,6 +64,17 @@ class ComputeTask:
else:
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):
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
if model_name is not None:
self.params["model_name"] = model_name
else:
self.params["model_name"] = "v1-5-pruned-emaonly"
def display(self) -> str:
return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
src/aios_kernel/local_stability_node.py
+48 -19
View File
@@ -6,6 +6,8 @@ import logging
import base64
from PIL import Image
import requests
from typing import Tuple
from pathlib import Path
from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
from .compute_node import ComputeNode
@@ -26,12 +28,18 @@ class Local_Stability_ComputeNode(ComputeNode):
@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")
if os.getenv("LOCAL_STABILITY_URL") is None:
user_config.add_user_config(
"local_stability_url", "local stability url", False, None)
if os.getenv("TEXT2IMG_OUTPUT_DIR") is None:
home_dir = Path.home()
output_dir = Path.joinpath(home_dir, "text2img_output")
Path.mkdir(output_dir, exist_ok=True)
user_config.add_user_config(
"text2img_output_dir", "text2image output dir", True, output_dir)
if os.getenv("TEXT2IMG_DEFAULT_MODEL") is None:
user_config.add_user_config(
"text2img_default_model", "text2img default model", True, "v1-5-pruned-emaonly")
def __init__(self) -> None:
super().__init__()
@@ -73,6 +81,8 @@ class Local_Stability_ComputeNode(ComputeNode):
if self.output_dir is None:
self.output_dir = "./"
self.output_dir = os.path.abspath(self.output_dir)
self.start()
@@ -85,36 +95,53 @@ class Local_Stability_ComputeNode(ComputeNode):
async def remove_task(self, task_id: str):
pass
def _make_post_request(self, url, json) -> Tuple[str, requests.Response]:
try:
response = requests.post(url, json=json)
if response.status_code != 200:
return f'{response.status_code}, {response.json()}', None
return None, response
except Exception as e:
return f"{e}", None
def _run_task(self, task: ComputeTask):
task.state = ComputeTaskState.RUNNING
model_name = task.params["model_name"]
prompts = task.params["prompts"]
prompt = task.params["prompt"]
logging.info(f"call local stability {model_name} prompts: {prompts}")
logging.info(f"call local stability {model_name} prompts: {prompt}")
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:
err, resp = self._make_post_request(f'{self.url}/sdapi/v1/options', payload)
if err is not None:
task.state = ComputeTaskState.ERROR
logger.error(
f"set local stability model failed. err:{response.json()['errors']}")
err_msg = f"Set local stability model failed. err:{err}"
logger.error(err_msg)
task.error_str = err_msg
return None
logging.info(f"set local stability model {model_name} success")
payload = {
"prompt": prompts,
"prompt": prompt,
"steps": 20
}
response = requests.post(
url=f'{self.url}/sdapi/v1/txt2img', json=payload)
r = response.json()
err, resp = self._make_post_request(f'{self.url}/sdapi/v1/txt2img', payload)
if err is not None:
task.state = ComputeTaskState.ERROR
err_msg = f"Failed. err:{err}"
logger.error(err_msg)
task.error_str = err_msg
return None
r = resp.json()
print(len(r['images']))
for i in r['images']:
image = Image.open(io.BytesIO(
base64.b64decode(i.split(",", 1)[0])))
@@ -128,6 +155,8 @@ class Local_Stability_ComputeNode(ComputeNode):
return result
task.error_str = "Unknown error!"
task.state = ComputeTaskState.ERROR
return None
def start(self):
src/aios_kernel/open_ai_node.py
+2 -2
View File
@@ -106,14 +106,14 @@ class OpenAI_ComputeNode(ComputeNode):
logger.info(f"call openai {mode_name} prompts: {prompts}")
resp = openai.ChatCompletion.create(model=mode_name,
messages=prompts,
max_tokens=result_token,
#max_tokens=result_token,
temperature=0.7)
else:
logger.info(f"call openai {mode_name} prompts: {prompts} functions: {json.dumps(llm_inner_functions)}")
resp = openai.ChatCompletion.create(model=mode_name,
messages=prompts,
functions=llm_inner_functions,
max_tokens=result_token,
#max_tokens=result_token,
temperature=0.7) # TODO: add temperature to task params?
src/aios_kernel/workflow_env.py
+20 -2
View File
@@ -9,6 +9,7 @@ import logging
from typing import Optional
from .text_to_speech_function import TextToSpeechFunction
from .compute_kernel import ComputeKernel
from .environment import Environment,EnvironmentEvent
from .ai_function import SimpleAIFunction
from .storage import AIStorage
@@ -82,6 +83,16 @@ class CalenderEnvironment(Environment):
"update event in calender",
self._update_event,update_param))
#maybe this function should be in other env?
paint_param = {
"prompt": "A description of the content of the painting",
"model_name": "Which model to use to draw the picture, can be None"
}
self.add_ai_function(SimpleAIFunction("paint",
"Draw a picture according to the description",
self._paint,paint_param))
#self.add_ai_function(SimpleAIFunction("user_confirm",
# "user confirm",
# self._user_confirm))
@@ -247,6 +258,14 @@ class CalenderEnvironment(Environment):
formatted_time = now.strftime('%Y-%m-%d %H:%M:%S')
return formatted_time
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}"
else:
return f'exec paint OK, saved as a local file, path is: {result.result["file"]}'
# Default Workflow Environment(Context)
class WorkflowEnvironment(Environment):
def __init__(self, env_id: str,db_file:str) -> None:
@@ -327,5 +346,4 @@ class WorkflowEnvironment(Environment):
logging.error(f"Error occurred while update env{self.env_id}.{key} ,error:{e}")
def get_functions(self):
pass
pass
src/service/aios_shell/aios_shell.py
+9
View File
@@ -23,6 +23,7 @@ directory = os.path.dirname(__file__)
sys.path.append(directory + '/../../')
from aios_kernel import AIOS_Version,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode,Local_Stability_ComputeNode
import proxy
from aios_kernel import *
@@ -62,6 +63,8 @@ class AIOS_Shell:
google_text_to_speech = GoogleTextToSpeechNode.get_instance()
google_text_to_speech.declare_user_config()
Local_Stability_ComputeNode.declare_user_config()
async def _handle_no_target_msg(self,bus:AIBus,msg:AgentMsg) -> bool:
target_id = msg.target.split(".")[0]
@@ -114,6 +117,12 @@ class AIOS_Shell:
await llama_ai_node.start()
# ComputeKernel.get_instance().add_compute_node(llama_ai_node)
local_sd_node = Local_Stability_ComputeNode.get_instance()
if await local_sd_node.initial() is not True:
logger.error("local stability node initial failed!")
return False
ComputeKernel.get_instance().add_compute_node(local_sd_node)
await ComputeKernel.get_instance().start()
AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg)