update

examples/sd_tool_usage.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+# @Date    : 1/11/2024 7:06 PM
+# @Author  : stellahong (stellahong@fuzhi.ai)
+# @Desc    :
+import asyncio
+from metagpt.const import METAGPT_ROOT
+from metagpt.actions.write_analysis_code import WriteCodeWithTools
+from metagpt.plan.planner import Planner
+from metagpt.actions.execute_code import ExecutePyCode
+from metagpt.roles.code_interpreter import CodeInterpreter
+
+sd_url = 'http://106.75.10.65:19094/sdapi/v1/txt2img'
+requirement = f"i have a text2image tool, generate a girl image use it, sd_url={sd_url}"
+
+if __name__ == "__main__":
+    code_interpreter = CodeInterpreter(use_tools=True, goal=requirement)
+    asyncio.run(code_interpreter.run(requirement))
+    # planner = Planner(
+    #     goal="i have a sdt2i tool, generate a girl image use it, sd_url='http://106.75.10.65:19094/sdapi/v1/txt2img'",
+    #     auto_run=True)
+    # asyncio.run(planner.update_plan())
+
+# schema_path = METAGPT_ROOT / "metagpt/tools/functions/schemas"
+# #
+# prompt = "1girl,  beautiful"
+# planner = Planner(
+#     goal="i have a sdt2i tool, generate a girl image use it, sd_url='http://106.75.10.65:19094/sdapi/v1/txt2img'",
+#     auto_run=True)
+# asyncio.run(planner.update_plan())
+# planner.plan.current_task.task_type = "sd"
+# planner.plan.current_task.instruction = "Use the sdt2i tool with the provided API endpoint to generate the girl image."
+# executor = ExecutePyCode()
+#
+# tool_context, code = asyncio.run(WriteCodeWithTools(schema_path=schema_path).run(
+#                     context=f"task prompt: {prompt}",
+#                     plan=planner.plan,
+#                     column_info="",
+#                 ))
+# print(code)
+# asyncio.run(executor.run(code))

metagpt/prompts/ml_engineer.py

@@ -58,6 +58,7 @@ Please assign a task type to each task in the list below from the given categori
 - **data_preprocess**: Only for changing value inplace.
 - **model_train**: Only for training model.
 - **model_evaluate**: Only for evaluating model.
+- **stable_diffusion**: Related to text2image, image2image using stable diffusion model.
 - **other**: Any tasks that do not fit into the previous categories, such as visualization, summarizing findings, etc.
 """
 
@@ -326,4 +327,5 @@ TASK_MODULE_MAP = {
     "data_preprocess": "metagpt.tools.functions.libs.data_preprocess",
     "feature_engineering": "metagpt.tools.functions.libs.feature_engineering",
     "udf": "metagpt.tools.functions.libs.udf",
+    "stable_diffusion": "metagpt.tools.sd_engine",
 }