ZIT prompting guide and tweaks

README.md

@@ -5,6 +5,11 @@ A Model Context Protocol (MCP) server designed to provide LLMs with efficient, t
 ## ⚠️ AI SLOP DISCLAIMER
 This entire project was vibe-coded by an AI. It is 100% slop code. Use it at your own risk.
 
+## 🚨 SECURITY WARNING
+**This server is designed to be run on `localhost` ONLY.** 
+It contains tools (such as `read_image` and `list_directory`) that allow the LLM to read arbitrary files from your filesystem. If you expose this server to the network or a public IP, any user or compromised AI could potentially read sensitive system files (e.g., SSH keys, `/etc/passwd`). 
+**NEVER run this server on a public-facing IP without implementing strict path validation.**
+
 ## Tools Overview
 
 ### 📁 `list_directory`
@@ -64,9 +69,10 @@ This server requires Python 3.10+ and the following packages:
 *   `starlette`: Lightweight ASGI framework.
 *   `Pillow`: Image processing and thumbnail generation.
 *   `requests`: For communicating with the Stable Diffusion API.
+*   `httpx`: For asynchronous API requests (e.g., Wikipedia).
 
 ```bash
-pip install uvicorn starlette Pillow requests
+pip install -r requirements.txt
 ```
 
 ### Setup
@@ -79,21 +85,38 @@ pip install uvicorn starlette Pillow requests
 
 ## Configuration (`config.py`)
 
-You can tune the server's behavior in `config.py`:
+Tuning the server's behavior is done via `config.py`.
 
-### Server & Logging
-- **`HOST` / `PORT`**: The network address and port the server binds to.
-- **`LOG_LEVEL`**: Logging verbosity (`DEBUG`, `INFO`, `WARNING`, `ERROR`).
-- **`LOG_FILE`**: Path to the server log file.
-- **`USER_AGENT`**: The User-Agent string used for API requests (e.g., Wikipedia). Use a browser-like string to avoid 403 Forbidden errors.
+### 🌐 Server & Logging
+| Parameter | Description | Default |
+| :--- | :--- | :--- |
+| `HOST` | The network address the server binds to. | `"127.0.0.1"` |
+| `PORT` | The port the server listens on. | `8000` |
+| `LOG_LEVEL` | Logging verbosity (`DEBUG`, `INFO`, `WARNING`, `ERROR`). | `"WARNING"` |
+| `LOG_FILE` | Absolute path to the server log file. | `ROOT_DIR / "debug.log"` |
+| `USER_AGENT` | User-Agent string for API requests (e.g., Wikipedia). | Browser-like string |
 
-### Model & Token Tuning
-- **`PATCH_SIZE`**: Set this to `(clip.vision.patch_size * n_merge)` for your specific model to ensure token-perfect resizing.
-- **`PREVIEW_TOKEN_BUDGET`**: Controls how many tokens the `preview_image` tool aims for (default: 70).
-- **`CONTACT_SHEET_COLS` / `ROWS`**: Adjust the grid size to fit within your model's maximum context window.
-- **`CONTACT_SHEET_THUMB_SIZE`**: The pixel size of thumbnails in the contact sheet.
+### 🎨 Stable Diffusion Integration
+| Parameter | Description | Default |
+| :--- | :--- | :--- |
+| `SD_URL` | Base URL of the SD WebUI/Forge instance. | `"http://127.0.0.1:7860"` |
+| `MODEL_PRESETS_PATH` | Path to the `model_presets.toml` file. | `ROOT_DIR / "model_presets.toml"` |
+| `RES_PRESETS_PATH` | Path to the `resolution_presets.toml` file. | `ROOT_DIR / "resolution_presets.toml"` |
+| `TAG_DATABASE_PATH` | Path to the Danbooru `tags.csv` file. | (Path to extension folder) |
+| `TAG_SEARCH_LIMIT` | Number of results returned by `search_tags` (direct or similar). | `20` |
 
-### Image & Font Settings
-- **`IMAGE_QUALITY`**: Adjust JPEG compression (1-100).
-- **`SYSTEM_FONT_NAMES`**: A list of font names Pillow should attempt to find in the system path.
-- **`FALLBACK_FONT_PATHS`**: Absolute paths to `.ttf` files used if no system fonts are found.
+### 🧠 Model & Token Tuning (Optimized for Gemma 4)
+| Parameter | Description | Default |
+| :--- | :--- | :--- |
+| `PATCH_SIZE` | Model's vision patch size in pixels. | `48` |
+| `PREVIEW_TOKEN_BUDGET` | Target token count for `preview_image` thumbnails. | `70` |
+| `CONTACT_SHEET_COLS` | Number of columns in the contact sheet grid. | `10` |
+| `CONTACT_SHEET_ROWS` | Number of rows in the contact sheet grid. | `7` |
+| `CONTACT_SHEET_THUMB_SIZE` | Pixel size of thumbnails in the contact sheet. | `192` |
+
+### 🖼️ Image & Font Settings
+| Parameter | Description | Default |
+| :--- | :--- | :--- |
+| `IMAGE_QUALITY` | JPEG compression quality (1-100). | `95` |
+| `SYSTEM_FONT_NAMES` | List of font names for Pillow to try in system paths. | Arial, DejaVu, etc. |
+| `FALLBACK_FONT_PATHS` | List of absolute paths to `.ttf` files. | Linux-specific paths |

config.py

@@ -1,5 +1,6 @@
 import logging
 from pathlib import Path
+from contextvars import ContextVar
 
 # --- Project Root ---
 # Get the directory where config.py is located
@@ -8,15 +9,20 @@ ROOT_DIR = Path(__file__).parent.resolve()
 # --- Server & Logs ---
 HOST = "127.0.0.1"
 PORT = 8000
-LOG_LEVEL = "WARNING" # Options: "DEBUG", "INFO", "WARNING", "ERROR"
+request_host = ContextVar("request_host", default=f"{HOST}:{PORT}")
+LOG_LEVEL = "DEBUG" # Options: "DEBUG", "INFO", "WARNING", "ERROR"
 LOG_FILE = str(ROOT_DIR / "debug.log")
-USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:151.0) Gecko/20100101 Firefox/151.0" # Stealth User-Agent to bypass Wikipedia's bot detection
+USER_AGENT = "MooseCP/1.0 Local MCP Server (https://long-cat.net/)"
 
 # --- Stable Diffusion Config ---
 SD_URL = "http://127.0.0.1:7860"
 MODEL_PRESETS_PATH = str(ROOT_DIR / "model_presets.toml")
 RES_PRESETS_PATH = str(ROOT_DIR / "resolution_presets.toml")
-TAG_DATABASE_PATH = "/home/matt/stable-diffusion-webui/extensions/a1111-sd-webui-tagcomplete/tags/danbooru.csv"
+TAG_DATABASE_PATH = str(ROOT_DIR / "danbooru.csv")
+TAG_SEARCH_LIMIT = 20
+ENABLE_TAG_WIKI = False
+DANBOORU_LOGIN = ""  # Your Danbooru username (Optional)
+DANBOORU_API_KEY = "" # Your Danbooru API key (Optional)
 
 # --- Model Specific Token Tuning (Tuned for Gemma 4) ---
 # Patch size is typically (clip.vision.patch_size * n_merge)

main.py

@@ -3,6 +3,7 @@ import logging
 import uvicorn
 import signal
 import os
+import sys
 from starlette.applications import Starlette
 from starlette.routing import Route
 from starlette.responses import Response, StreamingResponse, FileResponse
@@ -56,6 +57,10 @@ async def messages_endpoint(request):
     if not sid or sid not in mcp_logic.sessions:
         return Response("Session not found", status_code=404)
 
+    # Capture the host header to ensure image links work across the network
+    host_header = request.headers.get("host", f"{config.HOST}:{config.PORT}")
+    config.request_host.set(host_header)
+
     try:
         body = await request.json()
     except Exception as e:
@@ -116,7 +121,8 @@ app.add_middleware(
 )
 
 if __name__ == "__main__":
-    config = uvicorn.Config(
+    # Use a separate variable for uvicorn config to avoid shadowing the config module
+    uvicorn_config = uvicorn.Config(
         app=app, 
         host=config.HOST, 
         port=config.PORT, 
@@ -124,7 +130,7 @@ if __name__ == "__main__":
         timeout_graceful_shutdown=2 # Reduced from 5 to 2 seconds
     )
 
-    server = uvicorn.Server(config)
+    server = uvicorn.Server(uvicorn_config)
     
     def signal_handler(sig, frame):
         logger.info("Shutdown signal received")
@@ -140,6 +146,5 @@ if __name__ == "__main__":
     except KeyboardInterrupt:
         pass
     finally:
-        # Final fallback to ensure the process actually dies
-        # if the event loop is still hanging on a connection
-        os._exit(0)
+        logger.info("Server process exiting.")
+        sys.exit(0)

mcp_logic.py

@@ -42,23 +42,13 @@ class MCPServer:
                 )
             
             try:
-                result_data = await tool.handler(args)
+                content = await tool.handler(args)
             except ToolError as e:
                 # Convert tool errors into a text response so the LLM can understand and potentially fix the input
-                result_data = {"text": str(e)}
+                content = [{"type": "text", "text": str(e)}]
             except Exception as e:
                 # Catch-all for unexpected crashes to prevent server death
-                result_data = {"text": f"Unexpected internal error: {str(e)}"}
-            
-            # Handle results that are already lists of content (for multimodal/multi-part responses)
-            if isinstance(result_data, list):
-                content = result_data
-            elif isinstance(result_data, dict) and "text" in result_data:
-                content = [{"type": "text", "text": result_data["text"]}]
-            elif isinstance(result_data, dict) and result_data.get("type") == "image":
-                content = [result_data]
-            else:
-                content = [result_data]
+                content = [{"type": "text", "text": f"Unexpected internal error: {str(e)}"}]
             
             return {"content": content}
 

model_presets.toml

@@ -3,23 +3,25 @@
 # Description and Guide are mandatory. 
 # Other keys should be the Labels found in the /info endpoint.
 
-["noobaiXLNAIXL_vPred10Version"]
-description = "Anime-style model, very stylistic and not aesthetic-tuned. Can do any NSFW."
+["NoobAI XL"]
+description = "Anime-style model, very good at specific artist styles and character knowledge. Not aesthetic-tuned: needs precise, explicit prompting for best results. Can do any sort of NSFW."
 guide = """
-This model is not aesthetic tuned, it must be given explicit tags for everything. Omitted parts of the prompt will not default to something 'good'. This model is very fickle, you will almost always need to iterate on the prompt or resubmit to roll the best picture.
+This model is not aesthetic tuned, it must be given explicit tags for everything. Omitted parts of the prompt will not default to something 'good'. This model shows high variance, you can often get a very different image by just resubmitting the same prompt, so do not hesitate to try again.
 Has very excellent understanding of characters and artists down to extremely niche. Unless prompting an original character, their name is enough to decribe their appearance completely except for clothing.
-Accepts a list of comma-separated booru-style tags. Use spaces, not underscores, for tags.
+Accepts a list of comma-separated booru-style tags. Use spaces, not underscores, for tags. **Use the `search_tags` tool to verify your tags**.
 Prompts MUST follow this format: <1girl/1boy/1other/solo/couple/(can use multiple)>, <character(s)>, <series>, <artist>, <tags>
-Every prompt MUST include every one of the above sections.
+Every prompt MUST include every one of the above sections (the angle brackets are not part of the prompt).
 Quality tags such as "masterpiece", "best quality", "very awa", are a LAST RESORT, they override the artist tags. If absolutely required they should be prepended.
-It understands 'implicit' artist tags such as "official art" or "game cg" for the 'artist' immediately following the series.
+It understands every artist, so pick one appropriate for the image. If desired, it also knows 'implicit' artist tags such as "official art" or "game cg" for the 'artist' immediately following the series name.
 Natural language understanding is very limited, but can do things like "dark blue skirt" or natural language order of tags such as "lying, on bed".
 Do not use a negative prompt unless explicitly required to exclude something, your first prompt should have a blank negative prompt.
 Has the SDXL problem with hands, works best if hand posture is explicitly prompted.
 No default background, so prompts should include "outdoors", "indoors", or something like "patterned background" etc.
 CFG Scale from 3.0 - 5.5 but the default 4.0 is usually fine.
 """
+filename = "noobaiXLNAIXL_vPred10Version.safetensors"
 "Resolution Set" = "sdxl"
+preset = "xl"
 "CFG Scale" = 4.0
 "Hires. fix" = true
 "Denoising strength" = 0.5
@@ -31,3 +33,22 @@ CFG Scale from 3.0 - 5.5 but the default 4.0 is usually fine.
 "Sampling Method" = "Euler a"
 "Schedule Type" = "Uniform"
 "Rescale CFG" = 0.3
+
+["Z-Image-Turbo"]
+description = "General image generation model. Aesthetic tuned, gets good results first try. Can do softcore NSFW, e.g. underwear, breasts, asses. No full-frontal nudity."
+guide = """
+This model is aesthetic tuned, regenerating with the same prompt will yield essentially the same image. Change the prompt before resubmitting.
+It's a CFG 1.0 turbo model, the negative prompt has no effect.
+Understands natural language very well, uses Qwen 3 4B as the text encoder. The longer and more detailed prompt the better, take advantage of line breaks and formatting.
+If any part of the image is left out of the prompt, it will default to generic AI slop which is most NOT what you want. Be very explicit about each character's details. Appearance, ethnicity, age, individual outfit components, facial expression, pose, action, where they are looking, position in the image, etc. should ALL be included in the prompt. Characters can be referenced by naming them and using this name later in the prompt. This also helps the model avoid mixing traits between them.
+This also applies to the image itself. Composition, framing, lighting, image style, setting, background, etc. should all be explicitly specified in the prompt.
+Has some idiosyncrasies so you may need to iterate the prompt a few times to get around some weird artifacts. Think things like "green eyes" making them glow green, or "blush" making the entire face glow. Also really wants to make shirts tucked in for some reason. Examine the generated image closely and edit the prompt if needed.
+"""
+preset = "zit"
+filename = "z_image_turbo_bf16.safetensors"
+modules = ["ae.safetensors", "qwen_3_4b_abliterated.safetensors"]
+"Resolution Set" = "2k"
+"CFG Scale" = 1.0
+"Sampling Steps" = 9
+"Sampling Method" = "Euler"
+"Schedule Type" = "Beta"

requirements.txt

@@ -0,0 +1,5 @@
+uvicorn
+starlette
+requests
+httpx
+Pillow

resolution_presets.toml

@@ -7,3 +7,9 @@ widescreen = "1344x768"
 landscape = "1152x896"
 square = "1024x1024"
 portrait = "896x1152"
+
+[2k]
+widescreen = "2048x1152"
+landscape = "2048x1536"
+square = "2048x2048"
+portrait = "1536x2048"

tools/__init__.py

@@ -22,7 +22,7 @@ from .contact_sheet import handle as contact_sheet_handler
 from .list_directory import handle as list_directory_details_handler
 from .preview_image import handle as preview_image_handler
 from .get_text_context import handle as get_text_context_handler
-from .wikipedia import handle as wikipedia_handler
+from .browse_wikipedia import handle as wikipedia_handler
 from .get_model_info import handle as get_model_info_handler
 from .generate_image import handle as generate_image_handler
 from .search_tags import handle as search_tags_handler
@@ -148,16 +148,16 @@ TOOL_REGISTRY = [
                 "model_name": {"type": "string", "description": "The name of the model to use. Required."},
                 "prompt": {"type": "string", "description": "The prompt for the image. Required."},
                 "negative_prompt": {"type": "string", "description": "The negative prompt to exclude unwanted elements."},
-                "resolution_preset": {"type": "string", "description": "A named resolution preset (e.g., 'square', 'portrait'). Available options depend on the model."},
+                "resolution_preset": {"type": "string", "description": "A named resolution preset from the model info"},
                 "cfg_scale": {"type": "number", "description": "CFG scale for prompt adherence. Usually should be left omitted to select the default."},
             },
-            "required": ["model_name", "prompt"],
+            "required": ["model_name", "prompt", "resolution_preset"],
         },
         handler=generate_image_handler
     ),
     Tool(
         name="search_tags",
-        description="Searches the Danbooru tag database for tags matching a query. Returns the most popular tags including alias matches. Useful for finding the correct booru-style tags for anime models.",
+        description="Searches the Danbooru tag database for tags matching a query. Returns the Danbooru wiki page on an exact match. On any query, returns a list of similar tags as well as aliases. Useful for finding the correct booru-style tags for anime models or getting more information.",
         schema={
             "type": "object",
             "properties": {

tools/browse_wikipedia.py

@@ -1,9 +1,8 @@
-import urllib.request
-import urllib.parse
+import httpx
 import json
 import re
 import asyncio
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 from .utils import ToolError
 import config
 
@@ -13,22 +12,33 @@ def strip_html(text: str) -> str:
     """Removes HTML tags from a string using regex to provide clean text to the LLM."""
     return re.sub(r'<[^>]*>', '', text)
 
-def _make_request(params: Dict[str, Any]) -> Dict[str, Any]:
+def clean_wikitext(text: str) -> str:
     """
-    Synchronous helper to make the API request using urllib.
-    urllib is used instead of httpx to avoid TLS/HTTP fingerprinting 
-    that triggers 403 Forbidden responses from Wikipedia.
+    Removes the most distracting elements of raw Wikitext:
+    1. HTML comments (<!-- ... -->)
+    2. Citations (<ref /> and <ref>...</ref>)
     """
-    query_string = urllib.parse.urlencode(params)
-    url = f"{API_URL}?{query_string}"
+    # Remove HTML comments
+    text = re.sub(r'<!--.*?-->', '', text, flags=re.DOTALL)
+    # Remove self-closing citations FIRST to prevent them being seen as opening tags
+    text = re.sub(r'<ref[^>]*/>', '', text)
+    # Remove paired citations
+    text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL)
+    return text
 
+async def _make_request(params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Asynchronous helper to make the API request using httpx.
+    A custom User-Agent is used to avoid bot detection.
+    """
     headers = {
         "User-Agent": config.USER_AGENT
     }
     
-    req = urllib.request.Request(url, headers=headers)
-    with urllib.request.urlopen(req) as response:
-        return json.loads(response.read().decode('utf-8'))
+    async with httpx.AsyncClient(headers=headers, timeout=15.0) as client:
+        response = await client.get(API_URL, params=params)
+        response.raise_for_status()
+        return response.json()
 
 async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
     """
@@ -42,8 +52,8 @@ async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
         "format": "json",
         "srlimit": limit
     }
-    # Run synchronous urllib call in a thread to avoid blocking the event loop
-    data = await asyncio.to_thread(_make_request, params)
+    # Directly await the async request
+    data = await _make_request(params)
     
     search_results = data.get("query", {}).get("search", [])
     if not search_results:
@@ -66,10 +76,11 @@ async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
         
     return "\n".join(lines)
 
-async def _fetch_content(title: str, section_index: int) -> Optional[str]:
+async def _fetch_content(title: str, section_index: int) -> tuple[Optional[str], Optional[str]]:
     """
     Retrieves raw Wikitext for a specific section.
     section_index=0 returns the lead section.
+    Returns a tuple of (content, final_title).
     """
     params = {
         "action": "query",
@@ -80,27 +91,30 @@ async def _fetch_content(title: str, section_index: int) -> Optional[str]:
         "redirects": 1,
         "format": "json"
     }
-    data = await asyncio.to_thread(_make_request, params)
+    data = await _make_request(params)
     
     pages = data.get("query", {}).get("pages", {})
     if not pages:
-        return None
+        return None, None
         
     page_id = next(iter(pages))
     page = pages[page_id]
+    final_title = page.get("title")
     
     if "missing" in page:
-        return None
+        return None, None
         
     revisions = page.get("revisions", [])
     if not revisions:
-        return None
+        return None, final_title
         
-    return revisions[0].get("*")
+    content = revisions[0].get("*")
+    return (clean_wikitext(content) if content else None), final_title
 
-async def _fetch_toc(title: str) -> Optional[str]:
+async def _fetch_toc(title: str) -> tuple[Optional[str], Optional[str]]:
     """
     Retrieves the Table of Contents data and formats it hierarchically.
+    Returns a tuple of (toc, final_title).
     """
     params = {
         "action": "parse",
@@ -109,18 +123,19 @@ async def _fetch_toc(title: str) -> Optional[str]:
         "format": "json",
         "redirects": 1
     }
-    data = await asyncio.to_thread(_make_request, params)
+    data = await _make_request(params)
     
     parse_data = data.get("parse")
     if not parse_data:
-        return None
+        return None, None
         
+    final_title = parse_data.get("title")
     toc_data = parse_data.get("tocdata", {})
     sections = toc_data.get("sections", [])
     if not sections:
-        return "No table of contents found for this page."
+        return "No table of contents found for this page.", final_title
     
-    lines = [f"Table of Contents for \"{parse_data.get('title', title)}\":"]
+    lines = [f"Table of Contents for \"{final_title}\":"]
     for s in sections:
         level = s.get("tocLevel", 1)
         index = s.get("index")
@@ -128,9 +143,9 @@ async def _fetch_toc(title: str) -> Optional[str]:
         indent = "  " * (level - 1)
         lines.append(f"{indent}[{index}] {line}")
         
-    return "\n".join(lines)
+    return "\n".join(lines), final_title
 
-async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
+async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Main handler for the browse_wikipedia tool.
     Supports modes: summary, toc, section, and search.
@@ -143,18 +158,24 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     section_index = args.get("section_index")
     search_limit = args.get("search_limit", 5)
     
+    final_title = None
+    
     if mode == "search":
         result = await _search(title, search_limit, fallback=False)
     elif mode == "summary":
         # Lead section + ToC is the default 'summary' to guide the AI's next steps
-        content = await _fetch_content(title, 0)
+        content, title_from_content = await _fetch_content(title, 0)
+        final_title = title_from_content
         if content is None:
             result = await _search(title, search_limit, fallback=True)
         else:
-            toc = await _fetch_toc(title)
+            toc, title_from_toc = await _fetch_toc(title)
+            if title_from_toc:
+                final_title = title_from_toc
             result = f"{content}\n\n---\n\n{toc}"
     elif mode == "toc":
-        toc = await _fetch_toc(title)
+        toc, title_from_toc = await _fetch_toc(title)
+        final_title = title_from_toc
         if toc is None:
             result = await _search(title, search_limit, fallback=True)
         else:
@@ -162,7 +183,8 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     elif mode == "section":
         if section_index is None:
             raise ToolError("Missing required parameter 'section_index' for mode='section'")
-        content = await _fetch_content(title, int(section_index))
+        content, title_from_content = await _fetch_content(title, int(section_index))
+        final_title = title_from_content
         if content is None:
             result = await _search(title, search_limit, fallback=True)
         else:
@@ -170,4 +192,7 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     else:
         raise ToolError(f"Invalid mode '{mode}'. Supported modes: summary, toc, section, search")
     
-    return {"text": result}
+    if final_title and final_title != title:
+        result = f"Redirected to \"{final_title}\"\n\n{result}"
+            
+    return [{"type": "text", "text": result}]

tools/contact_sheet.py

@@ -2,7 +2,9 @@ import base64
 import logging
 import io
 import datetime
+import math
 from pathlib import Path
+
 from typing import Any, Dict, List
 from PIL import Image, ImageDraw, ImageFont, ImageOps
 import config
@@ -14,8 +16,7 @@ logger = logging.getLogger("MooseCP")
 async def handle(args: Dict[str, Any]):
     """
     Generates a contact sheet of images.
-    Grid: 10 columns x 7 rows (70 images total).
-    Sized for optimal token usage (1120 tokens) on llama.cpp.
+    Sized for optimal token usage according to config values.
     """
     dir_path_str = args.get("path")
     page = int(args.get("page", 1))
@@ -44,10 +45,10 @@ async def handle(args: Dict[str, Any]):
     }
     sort_label = sort_labels.get(sort_by, "Name (alphabetical)")
 
-    # Fixed grid dimensions for token optimization
-    COLS = config.CONTACT_SHEET_COLS
-    ROWS = config.CONTACT_SHEET_ROWS
-    PAGE_SIZE = COLS * ROWS
+    # Grid dimensions for token optimization (see config.py)
+    MAX_COLS = config.CONTACT_SHEET_COLS
+    MAX_ROWS = config.CONTACT_SHEET_ROWS
+    PAGE_SIZE = MAX_COLS * MAX_ROWS
 
     
     total_files = len(file_info_list)
@@ -56,9 +57,21 @@ async def handle(args: Dict[str, Any]):
     if not paged_files:
         return {"text": f"No images found on page {page}."}
 
-    thumb_size = 192 # 192 / 48 = 4 patches per side
-    canvas_w = COLS * thumb_size
-    canvas_h = ROWS * thumb_size
+    # Dynamic grid sizing to avoid blank space
+    n_images = len(paged_files)
+    thumb_size = config.CONTACT_SHEET_THUMB_SIZE
+    
+    if n_images <= config.CONTACT_SHEET_ROWS ** 2:
+        # Smallest square-ish grid that fits all
+        cols = math.ceil(math.sqrt(n_images))
+        rows = math.ceil(n_images / cols)
+    else:
+        # Lock rows, expand columns
+        rows = config.CONTACT_SHEET_ROWS
+        cols = math.ceil(n_images / config.CONTACT_SHEET_ROWS)
+    
+    canvas_w = cols * thumb_size
+    canvas_h = rows * thumb_size
     
     canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30))
     draw = ImageDraw.Draw(canvas)
@@ -87,8 +100,8 @@ async def handle(args: Dict[str, Any]):
     start_idx = (page - 1) * PAGE_SIZE
     for i, info in enumerate(paged_files):
         full_path = info["path"]
-        row = i // COLS
-        col = i % COLS
+        row = i // cols
+        col = i % cols
         
         x = col * thumb_size
         y = row * thumb_size

tools/generate_image.py

@@ -1,9 +1,60 @@
 import requests
 import base64
+import asyncio
 import config
+
 from typing import Any, Dict, List
 from tools.utils import ToolError, load_toml
 
+async def ensure_model_state(model_name: str, model_preset: Dict[str, Any]):
+    """
+    Checks the current server state and switches model/modules if they differ from the preset.
+    """
+    try:
+        config_resp = await asyncio.to_thread(requests.get, f"{config.SD_URL}/config", timeout=10)
+        config_resp.raise_for_status()
+        cfg_data = config_resp.json()
+        components = cfg_data.get("components", [])
+        
+        # Extract current state from components
+        current_state = {}
+        for comp in components:
+            elem_id = comp.get("props", {}).get("elem_id")
+            if elem_id:
+                current_state[elem_id] = comp.get("props", {}).get("value")
+        
+        # 1. Check and change Checkpoint
+        active_ckpt = current_state.get("setting_sd_model_checkpoint")
+        # Use 'filename' from TOML if available, otherwise fall back to the model_name key
+        target_ckpt = model_preset.get("filename", model_name)
+        
+        if active_ckpt != target_ckpt:
+            preset = model_preset.get("preset", "xl")
+            await asyncio.to_thread(
+                requests.post,
+                f"{config.SD_URL}/api/checkpoint_change", 
+                json={"data": [target_ckpt, preset]}, 
+                timeout=30
+            )
+        
+        # 2. Check and change VAE / Text Encoders
+        active_modules = current_state.get("setting_sd_modules", [])
+        target_modules = model_preset.get("modules", [])
+        
+        if active_modules != target_modules:
+            preset = model_preset.get("preset", "xl")
+            await asyncio.to_thread(
+                requests.post,
+                f"{config.SD_URL}/api/modules_change", 
+                json={"data": [target_modules, preset]}, 
+                timeout=30
+            )
+            
+    except Exception as e:
+        # We log this but don't necessarily raise a ToolError unless the generation itself fails,
+        # as the server might still be able to generate if it's just a state-check failure.
+        print(f"Warning: Failed to sync model state: {e}")
+
 async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Generates an image using the specified model and parameters.
@@ -17,25 +68,60 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     if not model_name:
         raise ToolError("The 'model_name' argument is required. Use get_model_info to see available models.")
 
+    res_preset_name = args.get("resolution_preset")
+    if not res_preset_name:
+        raise ToolError("The 'resolution_preset' argument is required. Use get_model_info to see available resolutions for the chosen model.")
+
     # 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
     try:
-        info_resp = requests.get(f"{config.SD_URL}/info", timeout=10)
+        info_resp = await asyncio.to_thread(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"]
     except Exception as e:
         raise ToolError(f"Failed to connect to Stable Diffusion server: {str(e)}")
 
-    # Build the label-to-index map
+    # Build the label-to-index map and a sparse payload
     label_map = {}
     label_counts = {}
     
-    payload = [p["parameter_default"] for p in params_info]
+    # Find the maximum param index to determine payload size
+    max_param_idx = 0
+    for p in params_info:
+        name = p.get("parameter_name", "")
+        if name.startswith("param_"):
+            try:
+                idx = int(name.replace("param_", ""))
+                max_param_idx = max(max_param_idx, idx)
+            except ValueError:
+                pass
 
-    for idx, p in enumerate(params_info):
-        label = p["label"]
+    # Initialize payload with Nones (size is max_idx + 1)
+    payload = [None] * (max_param_idx + 1)
+    
+    for idx_in_list, p in enumerate(params_info):
+        name = p.get("parameter_name", "")
+        label = p.get("label")
+        default = p.get("parameter_default")
+        
+        # Determine the absolute index in the payload
+        if name == "id_task":
+            abs_idx = 0
+        elif name.startswith("param_"):
+            try:
+                abs_idx = int(name.replace("param_", ""))
+            except ValueError:
+                continue
+        else:
+            # Fallback for unexpected names, though unlikely
+            continue
+
+        # Set the default value at the absolute index
+        payload[abs_idx] = default
+        
+        # Handle label mapping for AI overrides
         if not label or label.startswith("parameter_"):
-            label = p["parameter_name"]
+            label = name
             
         if label in label_counts:
             label_counts[label] += 1
@@ -44,17 +130,20 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
             label_counts[label] = 0
             mapped_label = label
             
-        label_map[mapped_label] = idx
+        label_map[mapped_label] = abs_idx
 
     # 3. LOAD CONFIGS
     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())}")
+        raise ToolError(f"Model '{model_name}' not found in presets. Please call get_model_info to see available models and their correct names.")
     
     model_preset = models_cfg[model_name]
 
+    # Ensure server state (Checkpoint, VAE, etc.) matches the preset before generating
+    await ensure_model_state(model_name, model_preset)
+
     # 4. MERGE PIPELINE
     for key, value in model_preset.items():
         if key in ["description", "guide", "Resolution Set"]:
@@ -85,7 +174,7 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
                 except (ValueError, AttributeError):
                     raise ToolError(f"Invalid resolution format for preset '{res_preset_name}': {res_val_str}. Expected 'WidthxHeight'.")
             else:
-                raise ToolError(f"Resolution preset '{res_preset_name}' not found for this model. Available: {', '.join(res_set.keys())}")
+                raise ToolError(f"Resolution preset '{res_preset_name}' not found for this model. You MUST call get_model_info with this model_name to see the available resolution presets.")
         else:
             raise ToolError(f"No resolution set configured for model '{model_name}'.")
 
@@ -105,13 +194,13 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     if "Seed" in label_map:
         payload[label_map["Seed"]] = -1
 
-    for _ in range(7):
-        payload.insert(39, None)
+    # The "Magic Number" splice is no longer needed as gaps are 
+    # automatically filled by the sparse-to-dense mapping.
 
     # 6. EXECUTE GENERATION
     try:
         gen_payload = {"data": payload}
-        gen_resp = requests.post(f"{config.SD_URL}/api/txt2img", json=gen_payload, timeout=300)
+        gen_resp = await asyncio.to_thread(requests.post, f"{config.SD_URL}/api/txt2img", json=gen_payload, timeout=300)
         gen_resp.raise_for_status()
         res_data = gen_resp.json()
         
@@ -129,12 +218,12 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
             raise ToolError(f"Could not extract file path from SD URL: {sd_img_url}")
 
         # Construct proxy URL through our MCP server
-        # We strip leading slash from raw_path to avoid double slashes in the proxy URL if we want, 
-        # but the current endpoint handles it.
-        proxy_url = f"http://localhost:{config.PORT}/file{raw_path}"
+        # Use the captured request host to ensure links work across the network
+        host = config.request_host.get()
+        proxy_url = f"http://{host}/file{raw_path}"
         
         # Download the image and convert to base64 for the AI's vision
-        img_resp = requests.get(sd_img_url, timeout=30)
+        img_resp = await asyncio.to_thread(requests.get, sd_img_url, timeout=30)
         img_resp.raise_for_status()
         b64_data = base64.b64encode(img_resp.content).decode('utf-8')
         

tools/get_model_info.py

@@ -3,7 +3,7 @@ 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]:
+async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Returns information about available models or detailed guides for a specific model.
     """
@@ -20,13 +20,30 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
                 "name": name,
                 "description": data.get("description", "No description provided.")
             })
-        return {
+        return [
+            {
+                "type": "text", 
                 "text": "Available models:\n\n" + "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) + 
-                    "\n\nTo get a detailed prompting guide and available resolutions for a specific model, call this tool again with the 'model_name' argument."
+                        "\n\nBefore generating an image, you MUST call this tool again with your selected model as the 'model_name' argument."
             }
+        ]
     
     if model_name not in models:
-        raise ToolError(f"Model '{model_name}' not found in presets. Available models: {', '.join(models.keys())}")
+        # Return the full catalog if the specific model isn't found
+        catalog = []
+        for name, data in models.items():
+            catalog.append({
+                "name": name,
+                "description": data.get("description", "No description provided.")
+            })
+        return [
+            {
+                "type": "text", 
+                "text": f"Model '{model_name}' not found.\n\nAvailable models:\n\n" + 
+                        "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) + 
+                        "\n\nBefore generating an image, you MUST call this tool again with selected model as the 'model_name' argument."
+            }
+        ]
     
     model_data = models[model_name]
     res_set_name = model_data.get("Resolution Set")
@@ -47,7 +64,9 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
     preset_text = "\n".join([f"- {k}: {v}" for k, v in presets_to_show.items()])
     res_text = ", ".join(res_options)
     
-    return {
+    return [
+        {
+            "type": "text", 
             "text": (
                 f"Model: {model_name}\n"
                 f"Description: {description}\n\n"
@@ -57,3 +76,4 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
                 f"Use 'resolution_preset' in generate_image to choose one of these."
             )
         }
+    ]

tools/get_text_context.py

@@ -35,7 +35,7 @@ async def handle(args: Dict[str, Any]):
     # Case 1: No pattern and no lines provided -> Dump whole file
     if not pattern and not lines_str:
         output = [f"{i+1}: {line}" for i, line in enumerate(all_lines)]
-        return {"text": "".join(output)}
+        return [{"type": "text", "text": "".join(output)}]
 
     # Determine target line numbers (1-based)
     target_lines: Set[int] = set()
@@ -61,7 +61,7 @@ async def handle(args: Dict[str, Any]):
             raise ToolError(f"Error parsing line indices: {e}")
 
     if not target_lines:
-        return {"text": "No matching lines found."}
+        return [{"type": "text", "text": "No matching lines found."}]
 
     # Expand targets with context
     lines_to_show: Set[int] = set()
@@ -83,4 +83,4 @@ async def handle(args: Dict[str, Any]):
         output.append(f"{ln}: {all_lines[ln-1]}")
         last_line = ln
 
-    return {"text": "".join(output)}
+    return [{"type": "text", "text": "".join(output)}]

tools/list_directory.py

@@ -43,7 +43,7 @@ async def handle(args: Dict[str, Any]):
     paged_items = all_items[start_idx:end_idx]
 
     if not paged_items:
-        return {"text": f"No items found on page {page}."}
+        return [{"type": "text", "text": f"No items found on page {page}."}]
 
     output = []
     for item in paged_items:
@@ -61,4 +61,4 @@ async def handle(args: Dict[str, Any]):
 
     header = f"Listing of {path_str}\nPage {effective_page} of {total_pages} ({total_items} total items). Showing {start_idx+1}-{min(end_idx, total_items)}."
     
-    return {"text": f"{header}\n\n" + "\n".join(output)}
+    return [{"type": "text", "text": f"{header}\n\n" + "\n".join(output)}]

tools/read_image.py

@@ -25,6 +25,6 @@ async def handle(args: Dict[str, Any]):
     
     try:
         data = base64.b64encode(path.read_bytes()).decode("utf-8")
-        return {"type": "image", "data": data, "mimeType": mime_type}
+        return [{"type": "image", "data": data, "mimeType": mime_type}]
     except Exception as e:
         raise ToolError(f"Error reading image: {str(e)}")

tools/read_metadata.py

@@ -26,11 +26,11 @@ async def handle(args: Dict[str, Any]):
             # Pillow parses PNG text chunks into the .info dictionary
             metadata = img.info
             if not metadata:
-                return {"text": "No metadata found in this PNG."}
+                return [{"type": "text", "text": "No metadata found in this PNG."}]
             
             # Format as a simple key: value list
             output = "\n".join([f"{k}: {v}" for k, v in metadata.items()])
-            return {"text": f"PNG Metadata:\n{output}"}
+            return [{"type": "text", "text": f"PNG Metadata:\n{output}"}]
     except ToolError:
         raise
     except Exception as e:

tools/search_tags.py

@@ -1,23 +1,24 @@
 import csv
 import os
-from typing import List, Dict, Any
+import config
+import difflib
+import httpx
+import re
+from typing import List, Dict, Any, Optional
 from tools.utils import ToolError, format_count, get_type_suffix
 
-# Tag database path is now managed in config.py
-async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Searches the Danbooru tag database for tags matching a query. 
-    Returns the most popular tags including alias matches.
-    """
-    query = args.get("query", "").lower()
-    if not query:
-        raise ToolError("The 'query' argument is required.")
+# Global cache to avoid reading the CSV from disk on every request
+_TAG_CACHE: Optional[List[Dict[str, Any]]] = None
 
+def _load_tags() -> List[Dict[str, Any]]:
+    """Reads the Danbooru tag CSV and caches it in memory."""
     tag_file_path = config.TAG_DATABASE_PATH
+    tags = []
+    
     if not os.path.exists(tag_file_path):
+        # We raise ToolError here because it's a fatal configuration issue
         raise ToolError(f"Tag database not found at {tag_file_path}. Please ensure the tag-autocomplete extension is installed.")
 
-    matches = []
     try:
         with open(tag_file_path, mode='r', encoding='utf-8') as f:
             reader = csv.reader(f)
@@ -25,44 +26,155 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
                 if not row or len(row) < 3:
                     continue
                 
-                name = row[0].lower()
+                name = row[0]
                 tag_type = row[1]
                 count = int(row[2]) if row[2].isdigit() else 0
                 aliases = row[3].lower().split(',') if len(row) > 3 else []
 
-                # Check for match in name or aliases
-                is_direct = query in name
-                is_alias = any(query in alias.strip() for alias in aliases)
-
-                if is_direct or is_alias:
-                    matches.append({
-                        "name": row[0],
+                tags.append({
+                    "name": name,
+                    "name_lower": name.lower(),
                     "type": tag_type,
                     "count": count,
-                        "matched_via": name if is_direct else "alias",
-                        "alias_match": "" if is_direct else next((a.strip() for a in aliases if query in a.strip()), query)
+                    "aliases": aliases
                 })
-
     except Exception as e:
         raise ToolError(f"Error reading tag database: {str(e)}")
     
-    # Sort by count descending
-    matches.sort(key=lambda x: x["count"], reverse=True)
+    return tags
 
-    # Format top 20 results
-    results = []
-    for m in matches[:20]:
-        count_fmt = format_count(str(m["count"]))
-        type_sfx = get_type_suffix(m["type"])
+    return tags
 
-        if m["matched_via"] == "alias":
-            line = f"{m['alias_match']} → {m['name']} ({count_fmt}){type_sfx}"
+async def _fetch_wiki_info(tag_name: str) -> Optional[str]:
+    """Fetches the wiki description for a tag from Danbooru."""
+    if not config.ENABLE_TAG_WIKI:
+        return None
+    
+    url = f"https://danbooru.donmai.us/wiki_pages/{tag_name}.json"
+    try:
+        # Use Basic Auth if credentials are provided
+        auth = None
+        if config.DANBOORU_LOGIN and config.DANBOORU_API_KEY:
+            auth = (config.DANBOORU_LOGIN, config.DANBOORU_API_KEY)
+
+        async with httpx.AsyncClient(timeout=5.0) as client:
+            # Using a browser-like UA to avoid potential blocks
+            headers = {"User-Agent": config.USER_AGENT}
+            resp = await client.get(url, headers=headers, auth=auth)
+            
+            if resp.status_code == 403:
+                return " (Wiki access blocked by Danbooru/Cloudflare)"
+            
+            if resp.status_code == 200:
+                data = resp.json()
+                wiki_body = data.get("wiki_page", {}).get("body")
+                if wiki_body:
+                    # Strip HTML tags for the LLM
+                    return re.sub(r'<[^>]*>', '', wiki_body).strip()
+    except Exception as e:
+        return f" (Error fetching wiki: {str(e)})"
+    return None
+
+async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """
+    Searches the Danbooru tag database for tags matching a query. 
+    Returns a unified list prioritized by substring matches then similarity.
+    """
+    global _TAG_CACHE
+    
+    # Normalize query: treat spaces and underscores as identical
+    query = args.get("query", "").lower().replace(" ", "_")
+    if not query:
+        raise ToolError("The 'query' argument is required.")
+
+    # Lazy-load the tags into memory
+    if _TAG_CACHE is None:
+        _TAG_CACHE = _load_tags()
+
+    # 1. Find Substring/Alias Matches (High Priority)
+    substring_matches = []
+    for tag in _TAG_CACHE:
+        is_direct = query in tag["name_lower"]
+        is_alias = any(query in alias.strip() for alias in tag["aliases"])
+        
+        if is_direct or is_alias:
+            # Find which alias actually matched for reporting
+            matched_alias = ""
+            if not is_direct:
+                matched_alias = next((a.strip() for a in tag["aliases"] if query in a.strip()), query)
+            
+            substring_matches.append({
+                "tag": tag,
+                "matched_via": "name" if is_direct else "alias",
+                "alias_match": matched_alias
+            })
+    
+    # Sort substring matches by count descending
+    substring_matches.sort(key=lambda x: x["tag"]["count"], reverse=True)
+
+    # 2. Find Similarity Matches (Low Priority)
+    all_names_lower = [t["name_lower"] for t in _TAG_CACHE]
+    similar_names_lower = difflib.get_close_matches(query, all_names_lower, n=config.TAG_SEARCH_LIMIT, cutoff=0.5)
+    similar_matches = []
+    for s_lower in similar_names_lower:
+        tag = next((t for t in _TAG_CACHE if t["name_lower"] == s_lower), None)
+        if tag:
+            similar_matches.append(tag)
+
+    # Build Unified List
+    final_results = []
+    
+    # Add substring matches first
+    for m in substring_matches:
+        final_results.append(m)
+        if len(final_results) >= config.TAG_SEARCH_LIMIT:
+            break
+            
+    # Fill remaining slots with similar matches
+    if len(final_results) < config.TAG_SEARCH_LIMIT:
+        mentioned_names = {m["tag"]["name_lower"] for m in final_results}
+        for tag in similar_matches:
+            if tag["name_lower"] not in mentioned_names:
+                final_results.append({"tag": tag, "matched_via": "similarity", "alias_match": ""})
+                if len(final_results) >= config.TAG_SEARCH_LIMIT:
+                    break
+
+    if not final_results:
+        return [{"type": "text", "text": f"No tags found matching '{query}'."}]
+
+    # Format output lines
+    output_lines = []
+    
+    # Special Case: Exact Match Wiki Header
+    # Check if the very first result is an exact match
+    first_res = final_results[0]
+    if first_res["tag"]["name_lower"] == query:
+        exact_tag = first_res["tag"]
+        count_fmt = format_count(str(exact_tag["count"]))
+        type_sfx = get_type_suffix(exact_tag["type"])
+        output_lines.append(f"Exact Match: {exact_tag['name']} ({count_fmt}){type_sfx}")
+        
+        wiki_info = await _fetch_wiki_info(exact_tag["name"])
+        if wiki_info:
+            output_lines.append(f"Wiki: {wiki_info}\n")
         else:
-            line = f"{m['name']} ({count_fmt}){type_sfx}"
+            output_lines.append("") # spacer
             
-        results.append(line)
+    # List the tags
+    for res in final_results:
+        tag = res["tag"]
+        count_fmt = format_count(str(tag["count"]))
+        type_sfx = get_type_suffix(tag["type"])
         
-    if not results:
-        return {"text": f"No tags found matching '{query}'."}
+        # If this is the exact match we already listed in the header, skip it
+        if first_res["tag"]["name_lower"] == query and tag["name_lower"] == query:
+            continue
 
-    return {"text": "Top matches:\n" + "\n".join(results)}
+        if res["matched_via"] == "alias":
+            line = f"- {res['alias_match']} → {tag['name']} ({count_fmt}){type_sfx}"
+        else:
+            line = f"- {tag['name']} ({count_fmt}){type_sfx}"
+        
+        output_lines.append(line)
+
+    return [{"type": "text", "text": "\n".join(output_lines)}]