Cleanup

tools/generate_image.py

@@ -1,23 +1,8 @@
 import requests
-import tomllib
-import os
 import base64
 import config
 from typing import Any, Dict, List
-from tools.utils import ToolError
-
-# Paths to config files
-MODEL_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "model_presets.toml")
-RES_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "resolution_presets.toml")
-
-SD_URL = "http://127.0.0.1:7860"
-
-def load_toml(path: str) -> Dict[str, Any]:
-    try:
-        with open(path, "rb") as f:
-            return tomllib.load(f)
-    except Exception as e:
-        raise ToolError(f"Failed to load config file {path}: {str(e)}")
+from tools.utils import ToolError, load_toml
 
 async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
@@ -34,7 +19,7 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
 
     # 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
     try:
-        info_resp = requests.get(f"{SD_URL}/info", timeout=10)
+        info_resp = requests.get(f"{config.SD_URL}/info", timeout=10)
         info_resp.raise_for_status()
         info_data = info_resp.json()
         params_info = info_data["named_endpoints"]["/txt2img"]["parameters"]
@@ -62,8 +47,8 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
         label_map[mapped_label] = idx
 
     # 3. LOAD CONFIGS
-    models_cfg = load_toml(MODEL_PRESETS_PATH)
-    res_cfg = load_toml(RES_PRESETS_PATH)
+    models_cfg = load_toml(config.MODEL_PRESETS_PATH)
+    res_cfg = load_toml(config.RES_PRESETS_PATH)
     
     if model_name not in models_cfg:
         raise ToolError(f"Model '{model_name}' not found in presets. Available: {', '.join(models_cfg.keys())}")