| @@ -8,6 +8,20 @@ from core.tools.provider.builtin.comfyui.tools.comfyui_client import ComfyUiClie | |||
| from core.tools.tool.builtin_tool import BuiltinTool | |||
| def sanitize_json_string(s): | |||
| escape_dict = { | |||
| "\n": "\\n", | |||
| "\r": "\\r", | |||
| "\t": "\\t", | |||
| "\b": "\\b", | |||
| "\f": "\\f", | |||
| } | |||
| for char, escaped in escape_dict.items(): | |||
| s = s.replace(char, escaped) | |||
| return s | |||
| class ComfyUIWorkflowTool(BuiltinTool): | |||
| def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMessage | list[ToolInvokeMessage]: | |||
| comfyui = ComfyUiClient(self.runtime.credentials["base_url"]) | |||
| @@ -26,13 +40,17 @@ class ComfyUIWorkflowTool(BuiltinTool): | |||
| set_prompt_with_ksampler = True | |||
| if "{{positive_prompt}}" in workflow: | |||
| set_prompt_with_ksampler = False | |||
| workflow = workflow.replace("{{positive_prompt}}", positive_prompt) | |||
| workflow = workflow.replace("{{negative_prompt}}", negative_prompt) | |||
| workflow = workflow.replace("{{positive_prompt}}", positive_prompt.replace('"', "'")) | |||
| workflow = workflow.replace("{{negative_prompt}}", negative_prompt.replace('"', "'")) | |||
| try: | |||
| prompt = json.loads(workflow) | |||
| except: | |||
| return self.create_text_message("the Workflow JSON is not correct") | |||
| except json.JSONDecodeError: | |||
| cleaned_string = sanitize_json_string(workflow) | |||
| try: | |||
| prompt = json.loads(cleaned_string) | |||
| except: | |||
| return self.create_text_message("the Workflow JSON is not correct") | |||
| if set_prompt_with_ksampler: | |||
| try: | |||