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())}")

tools/get_model_info.py

@@ -1,18 +1,7 @@
-import tomllib
-import os
-from typing import Any, Dict, Union, 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")
-
-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)}")
+import requests
+from typing import Any, Dict, List
+import config
+from tools.utils import ToolError, load_toml
 
 async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     """
@@ -20,8 +9,8 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     """
     model_name = args.get("model_name")
     
-    models = load_toml(MODEL_PRESETS_PATH)
-    res_presets = load_toml(RES_PRESETS_PATH)
+    models = load_toml(config.MODEL_PRESETS_PATH)
+    res_presets = load_toml(config.RES_PRESETS_PATH)
     
     if not model_name:
         # Return a catalog of all models

tools/utils.py

@@ -1,5 +1,6 @@
 import time
 import datetime
+import tomllib
 from pathlib import Path
 from typing import List, Dict, Any, Optional
 
@@ -8,6 +9,14 @@ class ToolError(Exception):
     """Custom exception for tool-related errors to be caught by the MCP server."""
     pass
 
+def load_toml(path: str) -> Dict[str, Any]:
+    """Loads a TOML file into a dictionary."""
+    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)}")
+
 def format_relative_time(timestamp: float) -> str:
     """Converts a timestamp to a human-readable relative format."""
     now = time.time()