mirror of
https://github.com/FoundationAgents/MetaGPT.git
synced 2026-06-29 15:59:42 +02:00
Merge branch 'code_intepreter' into code_intepreter_add_vision
# Conflicts: # metagpt/tools/__init__.py
This commit is contained in:
mannaandpoem
commit
23fdf90d21
19 changed files with 568 additions and 245 deletions
|
|
@ -85,20 +85,14 @@ class DebugCode(BaseWriteAnalysisCode):
|
|||
|
||||
async def run_reflection(
|
||||
self,
|
||||
# goal,
|
||||
# finished_code,
|
||||
# finished_code_result,
|
||||
context: List[Message],
|
||||
code,
|
||||
runtime_result,
|
||||
) -> dict:
|
||||
info = []
|
||||
# finished_code_and_result = finished_code + "\n [finished results]\n\n" + finished_code_result
|
||||
reflection_prompt = REFLECTION_PROMPT.format(
|
||||
debug_example=DEBUG_REFLECTION_EXAMPLE,
|
||||
context=context,
|
||||
# goal=goal,
|
||||
# finished_code=finished_code_and_result,
|
||||
code=code,
|
||||
runtime_result=runtime_result,
|
||||
)
|
||||
|
|
@ -106,33 +100,13 @@ class DebugCode(BaseWriteAnalysisCode):
|
|||
info.append(Message(role="system", content=system_prompt))
|
||||
info.append(Message(role="user", content=reflection_prompt))
|
||||
|
||||
# msg = messages_to_str(info)
|
||||
# resp = await self.llm.aask(msg=msg)
|
||||
resp = await self.llm.aask_code(messages=info, **create_func_config(CODE_REFLECTION))
|
||||
logger.info(f"reflection is {resp}")
|
||||
return resp
|
||||
|
||||
# async def rewrite_code(self, reflection: str = "", context: List[Message] = None) -> str:
|
||||
# """
|
||||
# 根据reflection重写代码
|
||||
# """
|
||||
# info = context
|
||||
# # info.append(Message(role="assistant", content=f"[code context]:{code_context}"
|
||||
# # f"finished code are executable, and you should based on the code to continue your current code debug and improvement"
|
||||
# # f"[reflection]: \n {reflection}"))
|
||||
# info.append(Message(role="assistant", content=f"[reflection]: \n {reflection}"))
|
||||
# info.append(Message(role="user", content=f"[improved impl]:\n Return in Python block"))
|
||||
# msg = messages_to_str(info)
|
||||
# resp = await self.llm.aask(msg=msg)
|
||||
# improv_code = CodeParser.parse_code(block=None, text=resp)
|
||||
# return improv_code
|
||||
|
||||
async def run(
|
||||
self,
|
||||
context: List[Message] = None,
|
||||
plan: str = "",
|
||||
# finished_code: str = "",
|
||||
# finished_code_result: str = "",
|
||||
code: str = "",
|
||||
runtime_result: str = "",
|
||||
) -> str:
|
||||
|
|
@ -140,14 +114,10 @@ class DebugCode(BaseWriteAnalysisCode):
|
|||
根据当前运行代码和报错信息进行reflection和纠错
|
||||
"""
|
||||
reflection = await self.run_reflection(
|
||||
# plan,
|
||||
# finished_code=finished_code,
|
||||
# finished_code_result=finished_code_result,
|
||||
code=code,
|
||||
context=context,
|
||||
runtime_result=runtime_result,
|
||||
)
|
||||
# 根据reflection结果重写代码
|
||||
# improv_code = await self.rewrite_code(reflection, context=context)
|
||||
improv_code = reflection["improved_impl"]
|
||||
return improv_code
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
@Author : orange-crow
|
||||
@File : code_executor.py
|
||||
"""
|
||||
import asyncio
|
||||
import re
|
||||
import traceback
|
||||
from abc import ABC, abstractmethod
|
||||
|
|
@ -81,6 +82,9 @@ class ExecutePyCode(ExecuteCode, Action):
|
|||
async def reset(self):
|
||||
"""reset NotebookClient"""
|
||||
await self.terminate()
|
||||
|
||||
# sleep 1s to wait for the kernel to be cleaned up completely
|
||||
await asyncio.sleep(1)
|
||||
await self.build()
|
||||
self.nb_client = NotebookClient(self.nb, timeout=self.timeout)
|
||||
|
||||
|
|
@ -181,7 +185,11 @@ class ExecutePyCode(ExecuteCode, Action):
|
|||
await self.nb_client.async_execute_cell(cell, cell_index)
|
||||
return True, ""
|
||||
except CellTimeoutError:
|
||||
return False, "TimeoutError"
|
||||
assert self.nb_client.km is not None
|
||||
await self.nb_client.km.interrupt_kernel()
|
||||
await asyncio.sleep(1)
|
||||
error_msg = "Cell execution timed out: Execution exceeded the time limit and was stopped; consider optimizing your code for better performance."
|
||||
return False, error_msg
|
||||
except DeadKernelError:
|
||||
await self.reset()
|
||||
return False, "DeadKernelError"
|
||||
|
|
|
|||
|
|
@ -60,7 +60,6 @@ class MLEngineer(CodeInterpreter):
|
|||
if code_execution_count > 0:
|
||||
logger.warning("We got a bug code, now start to debug...")
|
||||
code = await DebugCode().run(
|
||||
plan=self.planner.current_task.instruction,
|
||||
code=self.latest_code,
|
||||
runtime_result=self.working_memory.get(),
|
||||
context=self.debug_context,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@
|
|||
@File : __init__.py
|
||||
"""
|
||||
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -72,6 +71,12 @@ TOOL_TYPE_MAPPINGS = {
|
|||
desc="Only for evaluating model.",
|
||||
usage_prompt=MODEL_EVALUATE_PROMPT,
|
||||
),
|
||||
"stable_diffusion": ToolType(
|
||||
name="stable_diffusion",
|
||||
module="metagpt.tools.sd_engine",
|
||||
desc="Related to text2image, image2image using stable diffusion model.",
|
||||
usage_prompt="",
|
||||
),
|
||||
"vision": ToolType(
|
||||
name="vision",
|
||||
module=str(TOOL_LIBS_PATH / "vision"),
|
||||
|
|
|
|||
|
|
@ -37,8 +37,9 @@ class FillMissingValue(MLProcess):
|
|||
def transform(self, df: pd.DataFrame):
|
||||
if len(self.features) == 0:
|
||||
return df
|
||||
df[self.features] = self.si.transform(df[self.features])
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[self.features] = self.si.transform(new_df[self.features])
|
||||
return new_df
|
||||
|
||||
|
||||
class MinMaxScale(MLProcess):
|
||||
|
|
@ -54,8 +55,9 @@ class MinMaxScale(MLProcess):
|
|||
self.mms.fit(df[self.features])
|
||||
|
||||
def transform(self, df: pd.DataFrame):
|
||||
df[self.features] = self.mms.transform(df[self.features])
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[self.features] = self.mms.transform(new_df[self.features])
|
||||
return new_df
|
||||
|
||||
|
||||
class StandardScale(MLProcess):
|
||||
|
|
@ -71,8 +73,9 @@ class StandardScale(MLProcess):
|
|||
self.ss.fit(df[self.features])
|
||||
|
||||
def transform(self, df: pd.DataFrame):
|
||||
df[self.features] = self.ss.transform(df[self.features])
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[self.features] = self.ss.transform(new_df[self.features])
|
||||
return new_df
|
||||
|
||||
|
||||
class MaxAbsScale(MLProcess):
|
||||
|
|
@ -88,8 +91,9 @@ class MaxAbsScale(MLProcess):
|
|||
self.mas.fit(df[self.features])
|
||||
|
||||
def transform(self, df: pd.DataFrame):
|
||||
df[self.features] = self.mas.transform(df[self.features])
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[self.features] = self.mas.transform(new_df[self.features])
|
||||
return new_df
|
||||
|
||||
|
||||
class RobustScale(MLProcess):
|
||||
|
|
@ -105,8 +109,9 @@ class RobustScale(MLProcess):
|
|||
self.rs.fit(df[self.features])
|
||||
|
||||
def transform(self, df: pd.DataFrame):
|
||||
df[self.features] = self.rs.transform(df[self.features])
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[self.features] = self.rs.transform(new_df[self.features])
|
||||
return new_df
|
||||
|
||||
|
||||
class OrdinalEncode(MLProcess):
|
||||
|
|
@ -122,8 +127,9 @@ class OrdinalEncode(MLProcess):
|
|||
self.oe.fit(df[self.features])
|
||||
|
||||
def transform(self, df: pd.DataFrame):
|
||||
df[self.features] = self.oe.transform(df[self.features])
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[self.features] = self.oe.transform(new_df[self.features])
|
||||
return new_df
|
||||
|
||||
|
||||
class OneHotEncode(MLProcess):
|
||||
|
|
@ -142,9 +148,9 @@ class OneHotEncode(MLProcess):
|
|||
ts_data = self.ohe.transform(df[self.features])
|
||||
new_columns = self.ohe.get_feature_names_out(self.features)
|
||||
ts_data = pd.DataFrame(ts_data, columns=new_columns, index=df.index)
|
||||
df.drop(self.features, axis=1, inplace=True)
|
||||
df = pd.concat([df, ts_data], axis=1)
|
||||
return df
|
||||
new_df = df.drop(self.features, axis=1)
|
||||
new_df = pd.concat([new_df, ts_data], axis=1)
|
||||
return new_df
|
||||
|
||||
|
||||
class LabelEncode(MLProcess):
|
||||
|
|
@ -165,13 +171,14 @@ class LabelEncode(MLProcess):
|
|||
def transform(self, df: pd.DataFrame):
|
||||
if len(self.features) == 0:
|
||||
return df
|
||||
new_df = df.copy()
|
||||
for i in range(len(self.features)):
|
||||
data_list = df[self.features[i]].astype(str).tolist()
|
||||
for unique_item in np.unique(df[self.features[i]].astype(str)):
|
||||
if unique_item not in self.le_encoders[i].classes_:
|
||||
data_list = ["unknown" if x == unique_item else x for x in data_list]
|
||||
df[self.features[i]] = self.le_encoders[i].transform(data_list)
|
||||
return df
|
||||
new_df[self.features[i]] = self.le_encoders[i].transform(data_list)
|
||||
return new_df
|
||||
|
||||
|
||||
def get_column_info(df: pd.DataFrame) -> dict:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2023/11/17 10:33
|
||||
# @Author : lidanyang
|
||||
# @File : feature_engineering.py
|
||||
# @File : test_feature_engineering.py
|
||||
# @Desc : Feature Engineering Tools
|
||||
import itertools
|
||||
|
||||
|
|
@ -43,9 +43,9 @@ class PolynomialExpansion(MLProcess):
|
|||
ts_data = self.poly.transform(df[self.cols].fillna(0))
|
||||
column_name = self.poly.get_feature_names_out(self.cols)
|
||||
ts_data = pd.DataFrame(ts_data, index=df.index, columns=column_name)
|
||||
df.drop(self.cols, axis=1, inplace=True)
|
||||
df = pd.concat([df, ts_data], axis=1)
|
||||
return df
|
||||
new_df = df.drop(self.cols, axis=1)
|
||||
new_df = pd.concat([new_df, ts_data], axis=1)
|
||||
return new_df
|
||||
|
||||
|
||||
class CatCount(MLProcess):
|
||||
|
|
@ -57,8 +57,9 @@ class CatCount(MLProcess):
|
|||
self.encoder_dict = df[self.col].value_counts().to_dict()
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df[f"{self.col}_cnt"] = df[self.col].map(self.encoder_dict)
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[f"{self.col}_cnt"] = new_df[self.col].map(self.encoder_dict)
|
||||
return new_df
|
||||
|
||||
|
||||
class TargetMeanEncoder(MLProcess):
|
||||
|
|
@ -71,8 +72,9 @@ class TargetMeanEncoder(MLProcess):
|
|||
self.encoder_dict = df.groupby(self.col)[self.label].mean().to_dict()
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df[f"{self.col}_target_mean"] = df[self.col].map(self.encoder_dict)
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[f"{self.col}_target_mean"] = new_df[self.col].map(self.encoder_dict)
|
||||
return new_df
|
||||
|
||||
|
||||
class KFoldTargetMeanEncoder(MLProcess):
|
||||
|
|
@ -96,8 +98,9 @@ class KFoldTargetMeanEncoder(MLProcess):
|
|||
self.encoder_dict = tmp.groupby(self.col)[col_name].mean().to_dict()
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df[f"{self.col}_kf_target_mean"] = df[self.col].map(self.encoder_dict)
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[f"{self.col}_kf_target_mean"] = new_df[self.col].map(self.encoder_dict)
|
||||
return new_df
|
||||
|
||||
|
||||
class CatCross(MLProcess):
|
||||
|
|
@ -124,14 +127,15 @@ class CatCross(MLProcess):
|
|||
self.combs_map = dict(res)
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
new_df = df.copy()
|
||||
for comb in self.combs:
|
||||
new_col = f"{comb[0]}_{comb[1]}"
|
||||
_map = self.combs_map[new_col]
|
||||
df[new_col] = pd.Series(zip(df[comb[0]], df[comb[1]])).map(_map)
|
||||
new_df[new_col] = pd.Series(zip(new_df[comb[0]], new_df[comb[1]])).map(_map)
|
||||
# set the unknown value to a new number
|
||||
df[new_col].fillna(max(_map.values()) + 1, inplace=True)
|
||||
df[new_col] = df[new_col].astype(int)
|
||||
return df
|
||||
new_df[new_col].fillna(max(_map.values()) + 1, inplace=True)
|
||||
new_df[new_col] = new_df[new_col].astype(int)
|
||||
return new_df
|
||||
|
||||
|
||||
class GroupStat(MLProcess):
|
||||
|
|
@ -149,12 +153,12 @@ class GroupStat(MLProcess):
|
|||
self.group_df = group_df
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df = df.merge(self.group_df, on=self.group_col, how="left")
|
||||
return df
|
||||
new_df = df.merge(self.group_df, on=self.group_col, how="left")
|
||||
return new_df
|
||||
|
||||
|
||||
class SplitBins(MLProcess):
|
||||
def __init__(self, cols: str, strategy: str = "quantile"):
|
||||
def __init__(self, cols: list, strategy: str = "quantile"):
|
||||
self.cols = cols
|
||||
self.strategy = strategy
|
||||
self.encoder = None
|
||||
|
|
@ -164,8 +168,9 @@ class SplitBins(MLProcess):
|
|||
self.encoder.fit(df[self.cols].fillna(0))
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df[self.cols] = self.encoder.transform(df[self.cols].fillna(0))
|
||||
return df
|
||||
new_df = df.copy()
|
||||
new_df[self.cols] = self.encoder.transform(new_df[self.cols].fillna(0))
|
||||
return new_df
|
||||
|
||||
|
||||
class ExtractTimeComps(MLProcess):
|
||||
|
|
@ -192,91 +197,8 @@ class ExtractTimeComps(MLProcess):
|
|||
time_comps_df["dayofweek"] = time_s.dt.dayofweek + 1
|
||||
if "is_weekend" in self.time_comps:
|
||||
time_comps_df["is_weekend"] = time_s.dt.dayofweek.isin([5, 6]).astype(int)
|
||||
df = pd.concat([df, time_comps_df], axis=1)
|
||||
return df
|
||||
|
||||
|
||||
# @registry.register("feature_engineering", FeShiftByTime)
|
||||
# def fe_shift_by_time(df, time_col, group_col, shift_col, periods, freq):
|
||||
# df[time_col] = pd.to_datetime(df[time_col])
|
||||
#
|
||||
# def shift_datetime(date, offset, unit):
|
||||
# if unit in ["year", "y", "Y"]:
|
||||
# return date + relativedelta(years=offset)
|
||||
# elif unit in ["month", "m", "M"]:
|
||||
# return date + relativedelta(months=offset)
|
||||
# elif unit in ["day", "d", "D"]:
|
||||
# return date + relativedelta(days=offset)
|
||||
# elif unit in ["week", "w", "W"]:
|
||||
# return date + relativedelta(weeks=offset)
|
||||
# elif unit in ["hour", "h", "H"]:
|
||||
# return date + relativedelta(hours=offset)
|
||||
# else:
|
||||
# return date
|
||||
#
|
||||
# def shift_by_time_on_key(
|
||||
# inner_df, time_col, group_col, shift_col, offset, unit, col_name
|
||||
# ):
|
||||
# inner_df = inner_df.drop_duplicates()
|
||||
# inner_df[time_col] = inner_df[time_col].map(
|
||||
# lambda x: shift_datetime(x, offset, unit)
|
||||
# )
|
||||
# inner_df = inner_df.groupby([time_col, group_col], as_index=False)[
|
||||
# shift_col
|
||||
# ].mean()
|
||||
# inner_df.rename(columns={shift_col: col_name}, inplace=True)
|
||||
# return inner_df
|
||||
#
|
||||
# shift_df = df[[time_col, group_col, shift_col]].copy()
|
||||
# for period in periods:
|
||||
# new_col_name = f"{group_col}_{shift_col}_lag_{period}_{freq}"
|
||||
# tmp = shift_by_time_on_key(
|
||||
# shift_df, time_col, group_col, shift_col, period, freq, new_col_name
|
||||
# )
|
||||
# df = df.merge(tmp, on=[time_col, group_col], how="left")
|
||||
#
|
||||
# return df
|
||||
#
|
||||
#
|
||||
# @registry.register("feature_engineering", FeRollingByTime)
|
||||
# def fe_rolling_by_time(df, time_col, group_col, rolling_col, periods, freq, agg_funcs):
|
||||
# df[time_col] = pd.to_datetime(df[time_col])
|
||||
#
|
||||
# def rolling_by_time_on_key(inner_df, offset, unit, agg_func, col_name):
|
||||
# time_freq = {
|
||||
# "Y": [365 * offset, "D"],
|
||||
# "M": [30 * offset, "D"],
|
||||
# "D": [offset, "D"],
|
||||
# "W": [7 * offset, "D"],
|
||||
# "H": [offset, "h"],
|
||||
# }
|
||||
#
|
||||
# if agg_func not in ["mean", "std", "max", "min", "median", "sum", "count"]:
|
||||
# raise ValueError(f"Invalid agg function: {agg_func}")
|
||||
#
|
||||
# rolling_feat = inner_df.rolling(
|
||||
# f"{time_freq[unit][0]}{time_freq[unit][1]}", closed="left"
|
||||
# )
|
||||
# rolling_feat = getattr(rolling_feat, agg_func)()
|
||||
# depth = df.columns.nlevels
|
||||
# rolling_feat = rolling_feat.stack(list(range(depth)))
|
||||
# rolling_feat.name = col_name
|
||||
# return rolling_feat
|
||||
#
|
||||
# rolling_df = df[[time_col, group_col, rolling_col]].copy()
|
||||
# for period in periods:
|
||||
# for func in agg_funcs:
|
||||
# new_col_name = f"{group_col}_{rolling_col}_rolling_{period}_{freq}_{func}"
|
||||
# tmp = pd.pivot_table(
|
||||
# rolling_df,
|
||||
# index=time_col,
|
||||
# values=rolling_col,
|
||||
# columns=group_col,
|
||||
# )
|
||||
# tmp = rolling_by_time_on_key(tmp, period, freq, func, new_col_name)
|
||||
# df = df.merge(tmp, on=[time_col, group_col], how="left")
|
||||
#
|
||||
# return df
|
||||
new_df = pd.concat([df, time_comps_df], axis=1)
|
||||
return new_df
|
||||
|
||||
|
||||
class GeneralSelection(MLProcess):
|
||||
|
|
@ -302,8 +224,8 @@ class GeneralSelection(MLProcess):
|
|||
self.feats = feats
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df = df[self.feats + [self.label_col]]
|
||||
return df
|
||||
new_df = df[self.feats + [self.label_col]]
|
||||
return new_df
|
||||
|
||||
|
||||
class TreeBasedSelection(MLProcess):
|
||||
|
|
@ -344,8 +266,8 @@ class TreeBasedSelection(MLProcess):
|
|||
self.feats.append(self.label_col)
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df = df[self.feats]
|
||||
return df
|
||||
new_df = df[self.feats]
|
||||
return new_df
|
||||
|
||||
|
||||
class VarianceBasedSelection(MLProcess):
|
||||
|
|
@ -364,5 +286,5 @@ class VarianceBasedSelection(MLProcess):
|
|||
self.feats.append(self.label_col)
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df = df[self.feats]
|
||||
return df
|
||||
new_df = df[self.feats]
|
||||
return new_df
|
||||
|
|
|
|||
58
metagpt/tools/functions/schemas/stable_diffusion.yml
Normal file
58
metagpt/tools/functions/schemas/stable_diffusion.yml
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
SDEngine:
|
||||
type: class
|
||||
description: "Generate image using stable diffusion model"
|
||||
methods:
|
||||
__init__:
|
||||
description: "Initialize the SDEngine instance."
|
||||
parameters:
|
||||
properties:
|
||||
sd_url:
|
||||
type: str
|
||||
description: "URL of the stable diffusion service."
|
||||
simple_run_t2i:
|
||||
description: "Run the stable diffusion API for multiple prompts, calling the stable diffusion API to generate images."
|
||||
parameters:
|
||||
properties:
|
||||
payload:
|
||||
type: dict
|
||||
description: "Dictionary of input parameters for the stable diffusion API."
|
||||
auto_save:
|
||||
type: bool
|
||||
description: "Save generated images automatically."
|
||||
required:
|
||||
- prompts
|
||||
run_t2i:
|
||||
type: async function
|
||||
description: "Run the stable diffusion API for multiple prompts, calling the stable diffusion API to generate images."
|
||||
parameters:
|
||||
properties:
|
||||
payloads:
|
||||
type: list
|
||||
description: "List of payload, each payload is a dictionary of input parameters for the stable diffusion API."
|
||||
required:
|
||||
- payloads
|
||||
construct_payload:
|
||||
description: "Modify and set the API parameters for image generation."
|
||||
parameters:
|
||||
properties:
|
||||
prompt:
|
||||
type: str
|
||||
description: "Text input for image generation."
|
||||
required:
|
||||
- prompt
|
||||
returns:
|
||||
payload:
|
||||
type: dict
|
||||
description: "Updated parameters for the stable diffusion API."
|
||||
save:
|
||||
description: "Save generated images to the output directory."
|
||||
parameters:
|
||||
properties:
|
||||
imgs:
|
||||
type: str
|
||||
description: "Generated images."
|
||||
save_name:
|
||||
type: str
|
||||
description: "Output image name. Default is empty."
|
||||
required:
|
||||
- imgs
|
||||
|
|
@ -2,13 +2,14 @@
|
|||
# @Date : 2023/7/19 16:28
|
||||
# @Author : stellahong (stellahong@deepwisdom.ai)
|
||||
# @Desc :
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from os.path import join
|
||||
from typing import List
|
||||
|
||||
import requests
|
||||
from aiohttp import ClientSession
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
|
|
@ -51,9 +52,9 @@ default_negative_prompt = "(easynegative:0.8),black, dark,Low resolution"
|
|||
|
||||
|
||||
class SDEngine:
|
||||
def __init__(self):
|
||||
def __init__(self, sd_url=""):
|
||||
# Initialize the SDEngine with configuration
|
||||
self.sd_url = CONFIG.get("SD_URL")
|
||||
self.sd_url = sd_url if sd_url else CONFIG.get("SD_URL")
|
||||
self.sd_t2i_url = f"{self.sd_url}{CONFIG.get('SD_T2I_API')}"
|
||||
# Define default payload settings for SD API
|
||||
self.payload = payload
|
||||
|
|
@ -69,25 +70,36 @@ class SDEngine:
|
|||
):
|
||||
# Configure the payload with provided inputs
|
||||
self.payload["prompt"] = prompt
|
||||
self.payload["negtive_prompt"] = negtive_prompt
|
||||
self.payload["negative_prompt"] = negtive_prompt
|
||||
self.payload["width"] = width
|
||||
self.payload["height"] = height
|
||||
self.payload["override_settings"]["sd_model_checkpoint"] = sd_model
|
||||
logger.info(f"call sd payload is {self.payload}")
|
||||
return self.payload
|
||||
|
||||
def _save(self, imgs, save_name=""):
|
||||
def save(self, imgs, save_name=""):
|
||||
save_dir = CONFIG.workspace_path / SD_OUTPUT_FILE_REPO
|
||||
if not save_dir.exists():
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
batch_decode_base64_to_image(imgs, str(save_dir), save_name=save_name)
|
||||
|
||||
async def run_t2i(self, prompts: List):
|
||||
def simple_run_t2i(self, payload: dict, auto_save: bool = True):
|
||||
with requests.Session() as session:
|
||||
logger.debug(self.sd_t2i_url)
|
||||
rsp = session.post(self.sd_t2i_url, json=payload, timeout=600)
|
||||
|
||||
results = rsp.json()["images"]
|
||||
if auto_save:
|
||||
save_name = hashlib.sha256(payload["prompt"][:10].encode()).hexdigest()[:6]
|
||||
self.save(results, save_name=f"output_{save_name}")
|
||||
return results
|
||||
|
||||
async def run_t2i(self, payloads: List):
|
||||
# Asynchronously run the SD API for multiple prompts
|
||||
session = ClientSession()
|
||||
for payload_idx, payload in enumerate(prompts):
|
||||
for payload_idx, payload in enumerate(payloads):
|
||||
results = await self.run(url=self.sd_t2i_url, payload=payload, session=session)
|
||||
self._save(results, save_name=f"output_{payload_idx}")
|
||||
self.save(results, save_name=f"output_{payload_idx}")
|
||||
await session.close()
|
||||
|
||||
async def run(self, url, payload, session):
|
||||
|
|
@ -121,13 +133,3 @@ def batch_decode_base64_to_image(imgs, save_dir="", save_name=""):
|
|||
for idx, _img in enumerate(imgs):
|
||||
save_name = join(save_dir, save_name)
|
||||
decode_base64_to_image(_img, save_name=save_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
engine = SDEngine()
|
||||
prompt = "pixel style, game design, a game interface should be minimalistic and intuitive with the score and high score displayed at the top. The snake and its food should be easily distinguishable. The game should have a simple color scheme, with a contrasting color for the snake and its food. Complete interface boundary"
|
||||
|
||||
engine.construct_payload(prompt)
|
||||
|
||||
event_loop = asyncio.get_event_loop()
|
||||
event_loop.run_until_complete(engine.run_t2i(prompt))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue