Merge pull request #63 from stellaHSR/main

add UI role

metagpt/actions/action.py

@@ -13,7 +13,7 @@ from metagpt.actions.action_output import ActionOutput
 from tenacity import retry, stop_after_attempt, wait_fixed
 from pydantic import BaseModel
 from metagpt.utils.common import OutputParser
-
+from metagpt.logs import logger
 
 class Action(ABC):
     def __init__(self, name: str = '', context=None, llm: LLM = None):
@@ -55,8 +55,10 @@ class Action(ABC):
             system_msgs = []
         system_msgs.append(self.prefix)
         content = await self.llm.aask(prompt, system_msgs)
+        logger.debug(content)
         output_class = ActionOutput.create_model_class(output_class_name, output_data_mapping)
         parsed_data = OutputParser.parse_data_with_mapping(content, output_data_mapping)
+        logger.debug(parsed_data)
         instruct_content = output_class(**parsed_data)
         return ActionOutput(content, instruct_content)
 

metagpt/actions/write_prd.py

@@ -59,6 +59,7 @@ ATTENTION: Use '##' to SPLIT SECTIONS, not '#'. AND '## <SECTION_NAME>' SHOULD W
 
 ## Requirement Pool: Provided as Python list[str, str], the parameters are requirement description, priority(P0/P1/P2), respectively, comply with PEP standards; no more than 5 requirements and consider to make its difficulty lower
 
+## UI Design draft: Provide as Plain text. Be simple. Describe the elements and functions, also provide a simple style description and layout description.
 ## Anything UNCLEAR: Provide as Plain text. Make clear here.
 """
 FORMAT_EXAMPLE = """
@@ -105,6 +106,9 @@ The product should be a ...
 ]
 ```
 
+## UI Design draft
+Give a basic function description, and a draft
+
 ## Anything UNCLEAR
 There are no unclear points.
 ---
@@ -117,6 +121,7 @@ OUTPUT_MAPPING = {
     "Competitive Quadrant Chart": (str, ...),
     "Requirement Analysis": (str, ...),
     "Requirement Pool": (List[Tuple[str, str]], ...),
+    "UI Design draft":(str, ...),
     "Anything UNCLEAR": (str, ...),
 }
 
@@ -136,5 +141,6 @@ class WritePRD(Action):
 
         prompt = PROMPT_TEMPLATE.format(requirements=requirements, search_information=info,
                                         format_example=FORMAT_EXAMPLE)
+        logger.debug(prompt)
         prd = await self._aask_v1(prompt, "prd", OUTPUT_MAPPING)
         return prd

metagpt/roles/role.py

@@ -114,6 +114,7 @@ class Role:
     def _set_state(self, state):
         """Update the current state."""
         self._rc.state = state
+        logger.debug(self._actions)
         self._rc.todo = self._actions[self._rc.state]
 
     def set_env(self, env: 'Environment'):
@@ -170,8 +171,11 @@ class Role:
         if not self._rc.env:
             return 0
         env_msgs = self._rc.env.memory.get()
+        
         observed = self._rc.env.memory.get_by_actions(self._rc.watch)
+        
         already_observed = self._rc.memory.get()
+        
         news: list[Message] = []
         for i in observed:
             if i in already_observed:

metagpt/tools/sd_engine.py

@@ -0,0 +1,135 @@
+# -*- coding: utf-8 -*-
+# @Date    : 2023/7/19 16:28
+# @Author  : stellahong (stellahong@fuzhi.ai)
+# @Desc    :
+import os
+import asyncio
+from os.path import join
+from typing import List
+import json
+import io
+import base64
+
+from aiohttp import ClientSession
+from PIL import Image, PngImagePlugin
+
+from metagpt.logs import logger
+from metagpt.config import Config
+from metagpt.const import WORKSPACE_ROOT
+
+config = Config()
+
+payload = {
+    "prompt": "",
+    "negative_prompt": "(easynegative:0.8),black, dark,Low resolution",
+    "override_settings": {
+        "sd_model_checkpoint": "galaxytimemachinesGTM_photoV20"
+    },
+    "seed": -1,
+    "batch_size": 1,
+    "n_iter": 1,
+    "steps": 20,
+    "cfg_scale": 7,
+    "width": 512,
+    "height": 768,
+    "restore_faces": False,
+    "tiling": False,
+    "do_not_save_samples": False,
+    "do_not_save_grid": False,
+    'enable_hr': False,
+    'hr_scale': 2,
+    'hr_upscaler': 'Latent',
+    'hr_second_pass_steps': 0,
+    'hr_resize_x': 0,
+    'hr_resize_y': 0,
+    'hr_upscale_to_x': 0,
+    'hr_upscale_to_y': 0,
+    'truncate_x': 0,
+    'truncate_y': 0,
+    'applied_old_hires_behavior_to': None,
+    "eta": None,
+    
+    "sampler_index": "DPM++ SDE Karras",
+    "alwayson_scripts": {}
+}
+
+default_negative_prompt = "(easynegative:0.8),black, dark,Low resolution"
+
+
+class SDEngine:
+    def __init__(self):
+        # Initialize the SDEngine with configuration
+        self.config = Config()
+        self.sd_url = self.config.get('SD_URL')
+        self.sd_t2i_url = f"{self.sd_url}{self.config.get('SD_T2I_API')}"
+        # Define default payload settings for SD API
+        self.payload = payload
+        logger.info(self.sd_t2i_url)
+    
+    def construct_payload(self, prompt, negtive_prompt=default_negative_prompt, width=512, height=512,
+                          sd_model="galaxytimemachinesGTM_photoV20"):
+        # Configure the payload with provided inputs
+        self.payload["prompt"] = prompt
+        self.payload["negtive_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=""):
+        save_dir = WORKSPACE_ROOT / "resources"/"SD_Output"
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+        batch_decode_base64_to_image(imgs, save_dir, save_name=save_name)
+     
+    async def run_t2i(self, prompts: List):
+        # Asynchronously run the SD API for multiple prompts
+        session = ClientSession()
+        for payload_idx, payload in enumerate(prompts):
+            results = await self.run(url=self.sd_t2i_url, payload=payload, session=session)
+            self._save(results, save_name=f"output_{payload_idx}")
+        await session.close()
+    
+    async def run(self, url, payload, session):
+        # Perform the HTTP POST request to the SD API
+        async with session.post(url, json=payload, timeout=600) as rsp:
+            data = await rsp.read()
+        
+        rsp_json = json.loads(data)
+        imgs = rsp_json['images']
+        logger.info(f"callback rsp json is {rsp_json.keys()}")
+        return imgs
+    
+    async def run_i2i(self):
+        # todo: 添加图生图接口调用
+        raise NotImplementedError
+    
+    async def run_sam(self):
+        # todo：添加SAM接口调用
+        raise NotImplementedError
+
+def decode_base64_to_image(img, save_name):
+    image = Image.open(io.BytesIO(base64.b64decode(img.split(",", 1)[0])))
+    pnginfo = PngImagePlugin.PngInfo()
+    logger.info(save_name)
+    image.save(f"{save_name}.png", pnginfo=pnginfo)
+    return pnginfo, image
+
+
+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__":
+    import asyncio
+    
+    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))