iLOVE2D's picture
Upload 2846 files
5374a2d verified
from typing import Optional
import os
from ...tool import Toolkit
from ...storage_handler import FileStorageHandler
from .image_generation_edit import FluxImageGenerationEditTool
from ..openrouter_image_tools.image_analysis import ImageAnalysisTool
class FluxImageGenerationToolkit(Toolkit):
"""
Flux toolkit combining generation tool and a common image analysis tool.
"""
def __init__(
self,
name: str = "FluxImageGenerationToolkit",
api_key: Optional[str] = None,
save_path: str = "./imgs",
storage_handler: Optional[FileStorageHandler] = None,
analysis_model: str = "openai/gpt-4o-mini",
):
tools = []
gen_tool = FluxImageGenerationEditTool(
api_key=api_key,
storage_handler=storage_handler,
base_path=save_path,
)
tools.append(gen_tool)
try:
resolved_key = os.getenv("OPENROUTER_API_KEY")
if resolved_key:
analysis_tool = ImageAnalysisTool(
api_key=resolved_key,
model=analysis_model,
)
tools.append(analysis_tool)
except Exception:
pass
super().__init__(name=name, tools=tools)
self.api_key = api_key
self.save_path = save_path
self.storage_handler = storage_handler