18 lines
617 B
Python
18 lines
617 B
Python
from ai.tool import ToolSet
|
|
|
|
from .generate_image import GenerateImageTool
|
|
from .generate_image_anime import GenerateImageAnimeTool
|
|
|
|
|
|
class ImageGenerationToolSet(ToolSet):
|
|
def __init__(self, replicate_token: str):
|
|
functions = [
|
|
GenerateImageTool(replicate_token),
|
|
GenerateImageAnimeTool(replicate_token)
|
|
]
|
|
with open("ai/tools/image_generation/prompt.md", "r") as f:
|
|
system_prompt = f.read()
|
|
super().__init__(functions=functions, system_prompt=system_prompt)
|
|
|
|
|
|
__all__ = ["GenerateImageTool", "GenerateImageAnimeTool", "ImageGenerationToolSet"]
|