Tag not found update

model_presets.toml

@@ -3,12 +3,12 @@
 # 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.
 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.
@@ -19,7 +19,9 @@ Has the SDXL problem with hands, works best if hand posture is explicitly prompt
 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,18 @@ 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. NSFW is quite limited."
+guide = """
+This model is aesthetic tuned, regenerating with the same prompt will yield essentially the same image. Change the prompt before resubmitting.
+Understands natural language very well. Characters can be described by naming them and using this name later in the prompt. Longer, detailed prompts work better.
+"""
+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"

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

@@ -151,7 +151,7 @@ TOOL_REGISTRY = [
                 "resolution_preset": {"type": "string", "description": "A named resolution preset (e.g., 'square', 'portrait'). Available options depend on the model."},
                 "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
     ),

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/predict/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/predict/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,9 +68,13 @@ 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"]
@@ -55,6 +110,9 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     
     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"]:
@@ -111,7 +169,7 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     # 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()
         
@@ -134,7 +192,7 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
         proxy_url = f"http://localhost:{config.PORT}/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

@@ -26,7 +26,18 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
         }
     
     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 {
+            "text": f"Model '{model_name}' not found.\n\nAvailable 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."
+        }
     
     model_data = models[model_name]
     res_set_name = model_data.get("Resolution Set")

tools/search_tags.py

@@ -1,23 +1,22 @@
 import csv
 import os
-from typing import List, Dict, Any
+import config
+import difflib
+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,27 +24,53 @@ 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 []
 
+                tags.append({
+                    "name": name,
+                    "name_lower": name.lower(),
+                    "type": tag_type,
+                    "count": count,
+                    "aliases": aliases
+                })
+    except Exception as e:
+        raise ToolError(f"Error reading tag database: {str(e)}")
+    
+    return tags
+
+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.
+    """
+    global _TAG_CACHE
+    
+    query = args.get("query", "").lower()
+    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()
+
+    matches = []
+    for tag in _TAG_CACHE:
         # Check for match in name or aliases
-                is_direct = query in name
-                is_alias = any(query in alias.strip() for alias in aliases)
+        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:
             matches.append({
-                        "name": row[0],
-                        "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)
+                "name": tag["name"],
+                "type": tag["type"],
+                "count": tag["count"],
+                "matched_via": "name" if is_direct else "alias",
+                "alias_match": "" if is_direct else next((a.strip() for a in tag["aliases"] if query in a.strip()), query)
             })
 
-    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)
 
@@ -63,6 +88,23 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
         results.append(line)
 
     if not results:
+        # Attempt to find similar tags using difflib
+        all_names_lower = [t["name_lower"] for t in _TAG_CACHE]
+        suggestions_lower = difflib.get_close_matches(query, all_names_lower, n=10, cutoff=0.5)
+        
+        if not suggestions_lower:
             return {"text": f"No tags found matching '{query}'."}
         
+        # Map lowercased suggestions back to original tag objects
+        suggestions = []
+        for s_lower in suggestions_lower:
+            # Find the first tag that matches this lowercased name
+            tag = next((t for t in _TAG_CACHE if t["name_lower"] == s_lower), None)
+            if tag:
+                count_fmt = format_count(str(tag["count"]))
+                type_sfx = get_type_suffix(tag["type"])
+                suggestions.append(f"{tag['name']} ({count_fmt}){type_sfx}")
+        
+        return {"text": f"No exact matches for '{query}'. Did you mean:\n" + "\n".join(suggestions)}
+
     return {"text": "Top matches:\n" + "\n".join(results)}