Compare commits
45
Commits
5af30f8f2d
..
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
627ab3aec6 | ||
|
|
74eaa2a3e2 | ||
|
|
747bf5e02b | ||
|
|
8c18d0d40e | ||
|
|
bad4a24426 | ||
|
|
3e523e9725 | ||
|
|
b107aff150 | ||
|
|
1109a33edb | ||
|
|
e67aebbd3d | ||
|
|
2d30bd2155 | ||
|
|
f6f021f2a3 | ||
|
|
05a2a5010c | ||
|
|
f423764645 | ||
|
|
c0c3a4c405 | ||
|
|
f0e5267828 | ||
|
|
baa48d7cec | ||
|
|
81fef17af2 | ||
|
|
aa59df60d9 | ||
|
|
e9ee89f8a3 | ||
|
|
f676a07a71 | ||
|
|
51a4708092 | ||
|
|
79c975dd28 | ||
|
|
62c06b6254 | ||
|
|
ee6465d994 | ||
|
|
c24a37652c | ||
|
|
8beb7fe2a7 | ||
|
|
689cd1244c | ||
|
|
ca33869007 | ||
|
|
f808f4d599 | ||
|
|
449a41742d | ||
|
|
d8a2a102e4 | ||
|
|
31dc984197 | ||
|
|
50c06cc134 | ||
|
|
b6a820aa21 | ||
|
|
c705d30484 | ||
|
|
f05a662414 | ||
|
|
05e70a7fc4 | ||
|
|
504794f9a8 | ||
|
|
03e47bad42 | ||
|
|
a2cc56200c | ||
|
|
77a88c093d | ||
|
|
da5fdc3e4f | ||
|
|
d3d0ab9c63 | ||
|
|
48132aa21c | ||
|
|
d87fcb85be |
@@ -0,0 +1,122 @@
|
||||
# MooseCP Image Server
|
||||
|
||||
A Model Context Protocol (MCP) server designed to provide LLMs with efficient, token-optimized visual access to image directories and AI generation capabilities. Instead of dumping full-resolution images (which waste tokens and cause context overflow), MooseCP provides a hierarchical workflow: **List $\rightarrow$ Scan $\rightarrow$ Preview $\rightarrow$ Inspect**.
|
||||
|
||||
## ⚠️ 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`
|
||||
Provides a simplified file-browser view of a directory.
|
||||
* **Best for**: Getting a sense of the files present and their basic metadata (size, date).
|
||||
* **Features**: Pagination and sorting by name, size, or modification date.
|
||||
|
||||
### 🖼️ `contact_sheet`
|
||||
Generates a high-density grid of thumbnails.
|
||||
* **Best for**: Quickly scanning hundreds of images to find a specific one or get a general "vibe" of a folder.
|
||||
* **Note**: This tool is paginated. You must iterate through pages to see all images in a folder.
|
||||
* **Workflow**: Use the indices shown on the contact sheet to call `preview_image`.
|
||||
|
||||
### 🔍 `preview_image`
|
||||
Provides medium-detail previews optimized for the model's token budget.
|
||||
* **Best for**: Comparing a few candidates, inspecting specific details, or selecting a "favorite" image.
|
||||
* **Efficiency**: Automatically calculates dimensions to fit the model's patch size (e.g., 48px patches for Gemma 4), ensuring maximum detail without wasting tokens on padding.
|
||||
* **Workflow**: Pass indices from the `contact_sheet` or a direct file path.
|
||||
|
||||
### 📖 `read_png_metadata`
|
||||
Extracts AI generation parameters from PNG files.
|
||||
* **Best for**: Retrieving prompts, seeds, and model hashes from AI-generated images.
|
||||
|
||||
### 📸 `read_image`
|
||||
Returns the full-resolution image.
|
||||
* **Best for**: Final confirmation or deep visual analysis where every pixel counts.
|
||||
|
||||
### 🎨 `generate_image`
|
||||
Triggers an image generation on a local Stable Diffusion WebUI Forge instance.
|
||||
* **Workflow**: Always call `get_model_info` first to determine the correct prompting style (e.g., tag-based vs. natural language).
|
||||
* **Features**: Supports model-specific presets, resolution presets, and standard parameter overrides.
|
||||
* **Output**: Returns a Base64 image for AI analysis and a proxy URL for direct embedding in the chat.
|
||||
|
||||
### ℹ️ `get_model_info`
|
||||
Provides the "manual" for available generation models.
|
||||
* **Best for**: Learning the prompting style, recommended settings, and available resolution presets for a specific model.
|
||||
* **Workflow**: Call without arguments to see the catalog; call with `model_name` for the detailed guide.
|
||||
|
||||
### 🏷️ `search_tags`
|
||||
Searches the Danbooru tag database for recognized tags and aliases.
|
||||
* **Best for**: Finding the correct booru-style tags, checking tag popularity, and resolving aliases (e.g., 'lesbian' → 'yuri').
|
||||
* **Workflow**: Provide a query string to get a list of the most popular matching tags.
|
||||
|
||||
### 🌐 `browse_wikipedia`
|
||||
Allows the model to browse Wikipedia using its API.
|
||||
* **Best for**: Quickly retrieving summaries, structural maps (ToC), or specific section content from Wikipedia without dumping the entire page.
|
||||
* **Workflow**: Use `mode='summary'` (default) to get an overview and a Table of Contents. Use `mode='section'` with a linear index from the ToC to dive into specific details.
|
||||
* **Features**: Returns raw Wikitext to save tokens, handles redirects, and automatically falls back to a search result list if a page is not found.
|
||||
|
||||
---
|
||||
|
||||
## Installation & Requirements
|
||||
|
||||
### Dependencies
|
||||
This server requires Python 3.10+ and the following packages:
|
||||
* `uvicorn`: ASGI server for the SSE transport.
|
||||
* `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 -r requirements.txt
|
||||
```
|
||||
|
||||
### Setup
|
||||
1. Clone this repository to your server.
|
||||
2. (Optional) Edit `config.py` to adjust the server port, log level, or token budgets if you are using a model other than Gemma 4.
|
||||
3. Run the server:
|
||||
```bash
|
||||
python main.py
|
||||
```
|
||||
|
||||
## Configuration (`config.py`)
|
||||
|
||||
Tuning the server's behavior is done via `config.py`.
|
||||
|
||||
### 🌐 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 |
|
||||
|
||||
### 🎨 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` |
|
||||
|
||||
### 🧠 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 |
|
||||
@@ -0,0 +1,65 @@
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from contextvars import ContextVar
|
||||
|
||||
# --- Project Root ---
|
||||
# Get the directory where config.py is located
|
||||
ROOT_DIR = Path(__file__).parent.resolve()
|
||||
|
||||
# --- Server & Logs ---
|
||||
HOST = "127.0.0.1"
|
||||
PORT = 8000
|
||||
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 = "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 = 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)
|
||||
# For Gemma 4, this is 48px.
|
||||
PATCH_SIZE = 48
|
||||
|
||||
# The target token count for a single image preview.
|
||||
# The actual budget will be this or greater, but as small as possible.
|
||||
PREVIEW_TOKEN_BUDGET = 70
|
||||
|
||||
# Contact sheet grid optimized to fit within Gemma 4's 1120 max token budget
|
||||
# (10 cols * 7 rows) = 70 images.
|
||||
# Each thumb (192px) is 4x4 patches.
|
||||
# 70 images * (4*4) = 1120 tokens.
|
||||
CONTACT_SHEET_COLS = 10
|
||||
CONTACT_SHEET_ROWS = 7
|
||||
CONTACT_SHEET_THUMB_SIZE = 192
|
||||
|
||||
# --- Image Quality ---
|
||||
# JPEG image quality (usually max 95).
|
||||
# See: https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#jpeg-saving
|
||||
IMAGE_QUALITY = 95
|
||||
|
||||
# --- Font Configuration ---
|
||||
# Generic font names for cross-platform compatibility
|
||||
# Pillow's truetype() can often find these by name in the system path
|
||||
SYSTEM_FONT_NAMES = [
|
||||
"arialbd.ttf",
|
||||
"DejaVuSans-Bold",
|
||||
"LiberationSans-Bold",
|
||||
"Verdana",
|
||||
"Tahoma",
|
||||
]
|
||||
|
||||
# Fallback absolute paths for Linux systems
|
||||
FALLBACK_FONT_PATHS = [
|
||||
"/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
|
||||
"/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
|
||||
"/usr/share/fonts/truetype/freefont/FreeSansBold.ttf",
|
||||
]
|
||||
+140782
File diff suppressed because it is too large
Load Diff
@@ -3,22 +3,26 @@ 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
|
||||
from starlette.responses import Response, StreamingResponse, FileResponse
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
|
||||
from mcp_logic import MCPServer
|
||||
from tools import TOOL_REGISTRY
|
||||
import config
|
||||
|
||||
|
||||
# --- Configuration & Logging ---
|
||||
logging.basicConfig(
|
||||
level=logging.DEBUG,
|
||||
filename="debug.log",
|
||||
level=getattr(logging, config.LOG_LEVEL),
|
||||
filename=config.LOG_FILE,
|
||||
|
||||
filemode="a",
|
||||
format="%(asctime)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
logger = logging.getLogger("MattCP")
|
||||
logger = logging.getLogger("MooseCP")
|
||||
|
||||
mcp_logic = MCPServer()
|
||||
|
||||
@@ -53,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:
|
||||
@@ -72,10 +80,35 @@ async def messages_endpoint(request):
|
||||
await mcp_logic.send_to_session(sid, {"jsonrpc": "2.0", "id": req_id, "result": result})
|
||||
return Response(status_code=202)
|
||||
|
||||
async def file_endpoint(request):
|
||||
"""
|
||||
Proxy endpoint to serve files from disk with CORP/CORS headers to bypass browser restrictions.
|
||||
"""
|
||||
path = request.path_params.get("path")
|
||||
if not path:
|
||||
return Response("Path not provided", status_code=400)
|
||||
|
||||
# Ensure the path is absolute (SDFiles are usually in /tmp/gradio/...)
|
||||
# If the path doesn't start with /, we assume it's relative to root for simplicity in this context
|
||||
full_path = path if path.startswith("/") else f"/{path}"
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
return Response(f"File not found: {full_path}", status_code=404)
|
||||
|
||||
return FileResponse(
|
||||
full_path,
|
||||
headers={
|
||||
"Cross-Origin-Resource-Policy": "cross-origin",
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "GET"
|
||||
}
|
||||
)
|
||||
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route("/sse", endpoint=sse_endpoint),
|
||||
Route("/messages", endpoint=messages_endpoint, methods=["POST"]),
|
||||
Route("/file/{path:path}", endpoint=file_endpoint),
|
||||
]
|
||||
)
|
||||
|
||||
@@ -88,14 +121,16 @@ 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="127.0.0.1",
|
||||
port=8000,
|
||||
host=config.HOST,
|
||||
port=config.PORT,
|
||||
log_level="info",
|
||||
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")
|
||||
@@ -111,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)
|
||||
|
||||
+10
-12
@@ -3,6 +3,7 @@ import uuid
|
||||
import json
|
||||
from typing import Any, Dict, Union, List
|
||||
from starlette.responses import Response
|
||||
from tools.utils import ToolError
|
||||
|
||||
class MCPServer:
|
||||
def __init__(self):
|
||||
@@ -23,7 +24,7 @@ class MCPServer:
|
||||
return {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "MattCP", "version": "0.4.1"}
|
||||
"serverInfo": {"name": "MooseCP", "version": "0.4.1"}
|
||||
}
|
||||
|
||||
if method == "tools/list":
|
||||
@@ -40,17 +41,14 @@ class MCPServer:
|
||||
media_type="application/json"
|
||||
)
|
||||
|
||||
result_data = await tool.handler(args)
|
||||
|
||||
# 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]
|
||||
try:
|
||||
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
|
||||
content = [{"type": "text", "text": str(e)}]
|
||||
except Exception as e:
|
||||
# Catch-all for unexpected crashes to prevent server death
|
||||
content = [{"type": "text", "text": f"Unexpected internal error: {str(e)}"}]
|
||||
|
||||
return {"content": content}
|
||||
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# Model Presets
|
||||
# Format: ["Model Name"]
|
||||
# Description and Guide are mandatory.
|
||||
# Other keys should be the Labels found in the /info endpoint.
|
||||
|
||||
["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 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. **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 (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 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
|
||||
"Upscale by" = 1.5
|
||||
"Upscaler" = "Lanczos"
|
||||
"Hires steps" = 16
|
||||
"Hires CFG Scale" = 4.0
|
||||
"Sampling Steps" = 32
|
||||
"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"
|
||||
@@ -0,0 +1,5 @@
|
||||
uvicorn
|
||||
starlette
|
||||
requests
|
||||
httpx
|
||||
Pillow
|
||||
@@ -0,0 +1,15 @@
|
||||
# Resolution Presets
|
||||
# Format: [preset_name]
|
||||
# Values: "WidthxHeight"
|
||||
|
||||
[sdxl]
|
||||
widescreen = "1344x768"
|
||||
landscape = "1152x896"
|
||||
square = "1024x1024"
|
||||
portrait = "896x1152"
|
||||
|
||||
[2k]
|
||||
widescreen = "2048x1152"
|
||||
landscape = "2048x1536"
|
||||
square = "2048x2048"
|
||||
portrait = "1536x2048"
|
||||
@@ -1,5 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Start the MCP Image Server in the background
|
||||
nohup /home/matt/code/llama.cpp/build/bin/mcp-image-server/venv/bin/python3 /home/matt/code/llama.cpp/build/bin/mcp-image-server/server.py > /home/matt/code/llama.cpp/build/bin/mcp-image-server/server.log 2>&1 &
|
||||
echo "MCP Image Server started in background on port 8000"
|
||||
echo "Log file: /home/matt/code/llama.cpp/build/bin/mcp-image-server/server.log"
|
||||
+103
-9
@@ -18,14 +18,20 @@ class Tool:
|
||||
# Import handlers from separate files
|
||||
from .read_image import handle as read_image_handler
|
||||
from .read_metadata import handle as read_png_metadata_handler
|
||||
from .list_thumbnails import handle as list_thumbnails_handler
|
||||
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 .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
|
||||
|
||||
# Central registry of all available tools
|
||||
TOOL_REGISTRY = [
|
||||
Tool(
|
||||
name="read_image",
|
||||
description="Reads an image from the disk and returns it as an image content object.",
|
||||
description="Reads an image at full resolution for maximum detail. Use this for final confirmation of a selected image or when deep visual analysis of a specific file is required.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {"path": {"type": "string", "description": "Path to the image file"}},
|
||||
@@ -35,7 +41,7 @@ TOOL_REGISTRY = [
|
||||
),
|
||||
Tool(
|
||||
name="read_png_metadata",
|
||||
description="Reads the metadata (tEXt, zTXt, iTXt) from a PNG image to fetch info such as AI prompts.",
|
||||
description="Reads the metadata (tEXt, zTXt, iTXt) from a PNG image to fetch info such as prompts for AI generated images.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {"path": {"type": "string", "description": "Path to the PNG file"}},
|
||||
@@ -44,23 +50,22 @@ TOOL_REGISTRY = [
|
||||
handler=read_png_metadata_handler
|
||||
),
|
||||
Tool(
|
||||
name="list_thumbnails",
|
||||
description="Generates a thumbnail grid of images in a directory. Supports pagination and sorting.",
|
||||
name="contact_sheet",
|
||||
description="Generates a coarse, high-density overview of images in a directory. This tool is paginated; a single page may not show all images. To perform a comprehensive scan or find specific images, you MUST iterate through multiple pages. If ANYTHING requires selecting individual images, then use indices as inputs to preview_image. Rely on this tool ONLY for a general characterization of images in a directory, NOT anything about specific images.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Path to the directory containing images"},
|
||||
"page": {"type": "integer", "description": "The page number to retrieve (1-indexed)", "default": 1},
|
||||
"page_size": {"type": "integer", "description": "Number of images per page", "default": 64},
|
||||
"page": {"type": "integer", "description": "The page number to retrieve (1-indexed). Increment this value to navigate through the full list of images in the folder.", "default": 1},
|
||||
"sort_by": {"type": "string", "description": "Sort order: 'mtime' (newest first, default), 'name' (alphabetical), or 'size' (largest first)", "default": "mtime"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
handler=list_thumbnails_handler
|
||||
handler=contact_sheet_handler
|
||||
),
|
||||
Tool(
|
||||
name="list_directory",
|
||||
description="Lists the contents of a directory in a file browser format. Directories first, then files. Supports pagination and sorting.",
|
||||
description="Lists the contents of a directory in a simplified text format. Use this for checking file existence, sizes, or dates or for quick filesystem navigation. DO NOT use this tool for any sort of visual or image selection.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
@@ -73,4 +78,93 @@ TOOL_REGISTRY = [
|
||||
},
|
||||
handler=list_directory_details_handler
|
||||
),
|
||||
Tool(
|
||||
name="preview_image",
|
||||
description="Provides low-resolution previews and information such as filenames, size, and resolution for specific images. It uses the minimum viable token budget, so details WILL be missed. Only use this tool for quick down-selection of images from contact_sheet or where accuracy does not matter. For final confirmation or deep analysis of a single selected image, you MUST use read_image after this tool.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Path to an image file or a directory containing images"},
|
||||
"indices": {"type": "string", "description": "1-based indices or ranges (e.g., '1, 3-5, 10') to preview from the directory. Required if path is a directory."},
|
||||
"sort_by": {"type": "string", "description": "Sort order to resolve indices: 'mtime' (default), 'name', or 'size'."},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
handler=preview_image_handler
|
||||
),
|
||||
Tool(
|
||||
name="get_text_context",
|
||||
description="Extracts specific sections of a file with absolute line numbers and surrounding context. This tool MUST be called immediately before edit_file to verify the exact line numbers and content of the block being replaced, preventing misalignment errors.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Path to the file"},
|
||||
"pattern": {"type": "string", "description": "Regex pattern to search for"},
|
||||
"lines": {"type": "string", "description": "1-based indices or ranges (e.g., '10-20, 45')"},
|
||||
"context": {"type": "integer", "description": "Number of lines of context to show above and below", "default": 1},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
handler=get_text_context_handler
|
||||
),
|
||||
Tool(
|
||||
name="browse_wikipedia",
|
||||
description="Browse Wikipedia pages to retrieve information. Defaults to the page summary. Can also retrieve the table of contents or a specific section. If a page is not found, it automatically returns search results.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"title": {"type": "string", "description": "The title of the Wikipedia page to browse or the search query."},
|
||||
"mode": {
|
||||
"type": "string",
|
||||
"description": "The retrieval mode. 'summary' (default) for the lead section, 'toc' for the table of contents, 'section' for a specific section, or 'search' for explicit search results.",
|
||||
"enum": ["summary", "toc", "section", "search"],
|
||||
"default": "summary"
|
||||
},
|
||||
"section_index": {"type": "integer", "description": "The linear index of the section to retrieve. Required if mode='section'. This index can be found by calling the tool in 'toc' mode."},
|
||||
"search_limit": {"type": "integer", "description": "The maximum number of search results to return. Default is 5.", "default": 5},
|
||||
},
|
||||
"required": ["title"],
|
||||
},
|
||||
handler=wikipedia_handler
|
||||
),
|
||||
Tool(
|
||||
name="get_model_info",
|
||||
description="Returns a list of available txt2img models and their short descriptions. If a specific model name is provided, it returns the comprehensive prompting guide, available resolution presets, and active configuration tips for that model. CRITICAL: You must call this for any model you are unfamiliar with, as prompting styles vary wildly (e.g., tag-based vs. natural language) and using the wrong style will result in poor image quality.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"model_name": {"type": "string", "description": "The name of the model to get detailed info for. Leave empty to list all available models."}
|
||||
},
|
||||
"required": [],
|
||||
},
|
||||
handler=get_model_info_handler
|
||||
),
|
||||
Tool(
|
||||
name="generate_image",
|
||||
description="Generates an image using the specified model and parameters. To ensure high quality, verify the model's prompting requirements via get_model_info before calling this tool. You will usually not get the correct result the first time. If the generated image needs work, change the prompt and call this tool again. If it's good, use the provided markdown to display it to the user.",
|
||||
schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"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 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", "resolution_preset"],
|
||||
},
|
||||
handler=generate_image_handler
|
||||
),
|
||||
Tool(
|
||||
name="search_tags",
|
||||
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": {
|
||||
"query": {"type": "string", "description": "The tag or partial tag to search for (e.g., 'white hair' or 'lesbian')."}
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
handler=search_tags_handler
|
||||
),
|
||||
]
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
import httpx
|
||||
import json
|
||||
import re
|
||||
import asyncio
|
||||
from typing import Any, Dict, List, Optional
|
||||
from .utils import ToolError
|
||||
import config
|
||||
|
||||
API_URL = "https://en.wikipedia.org/w/api.php"
|
||||
|
||||
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 clean_wikitext(text: str) -> str:
|
||||
"""
|
||||
Removes the most distracting elements of raw Wikitext:
|
||||
1. HTML comments (<!-- ... -->)
|
||||
2. Citations (<ref /> and <ref>...</ref>)
|
||||
"""
|
||||
# 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
|
||||
}
|
||||
|
||||
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:
|
||||
"""
|
||||
Searches Wikipedia for a title.
|
||||
fallback=True: Used when a direct page request fails, providing 'Article not found' header.
|
||||
"""
|
||||
params = {
|
||||
"action": "query",
|
||||
"list": "search",
|
||||
"srsearch": title,
|
||||
"format": "json",
|
||||
"srlimit": limit
|
||||
}
|
||||
# Directly await the async request
|
||||
data = await _make_request(params)
|
||||
|
||||
search_results = data.get("query", {}).get("search", [])
|
||||
if not search_results:
|
||||
return f"No Wikipedia results found for '{title}'."
|
||||
|
||||
if fallback:
|
||||
header = f"Article not found. You MUST select one of the following articles and submit another query:"
|
||||
else:
|
||||
header = f"Search results for '{title}':"
|
||||
|
||||
lines = [header]
|
||||
for i, res in enumerate(search_results, 1):
|
||||
title_res = res.get("title")
|
||||
snippet = strip_html(res.get("snippet", ""))
|
||||
# Brackets around title help the AI identify the exact string for subsequent calls
|
||||
lines.append(f"{i}. [{title_res}] - Snippet: {snippet}")
|
||||
|
||||
if fallback:
|
||||
lines.append("\nUse the title in brackets [ ] to explore the selected page.")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
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",
|
||||
"prop": "revisions",
|
||||
"rvprop": "content",
|
||||
"rvsection": section_index,
|
||||
"titles": title,
|
||||
"redirects": 1,
|
||||
"format": "json"
|
||||
}
|
||||
data = await _make_request(params)
|
||||
|
||||
pages = data.get("query", {}).get("pages", {})
|
||||
if not pages:
|
||||
return None, None
|
||||
|
||||
page_id = next(iter(pages))
|
||||
page = pages[page_id]
|
||||
final_title = page.get("title")
|
||||
|
||||
if "missing" in page:
|
||||
return None, None
|
||||
|
||||
revisions = page.get("revisions", [])
|
||||
if not revisions:
|
||||
return None, final_title
|
||||
|
||||
content = revisions[0].get("*")
|
||||
return (clean_wikitext(content) if content else None), final_title
|
||||
|
||||
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",
|
||||
"page": title,
|
||||
"prop": "tocdata",
|
||||
"format": "json",
|
||||
"redirects": 1
|
||||
}
|
||||
data = await _make_request(params)
|
||||
|
||||
parse_data = data.get("parse")
|
||||
if not parse_data:
|
||||
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.", final_title
|
||||
|
||||
lines = [f"Table of Contents for \"{final_title}\":"]
|
||||
for s in sections:
|
||||
level = s.get("tocLevel", 1)
|
||||
index = s.get("index")
|
||||
line = s.get("line")
|
||||
indent = " " * (level - 1)
|
||||
lines.append(f"{indent}[{index}] {line}")
|
||||
|
||||
return "\n".join(lines), final_title
|
||||
|
||||
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.
|
||||
"""
|
||||
title = args.get("title")
|
||||
if not title:
|
||||
raise ToolError("Missing required parameter: 'title'")
|
||||
|
||||
mode = args.get("mode", "summary")
|
||||
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, 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, 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, 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:
|
||||
result = toc
|
||||
elif mode == "section":
|
||||
if section_index is None:
|
||||
raise ToolError("Missing required parameter 'section_index' for mode='section'")
|
||||
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:
|
||||
result = content
|
||||
else:
|
||||
raise ToolError(f"Invalid mode '{mode}'. Supported modes: summary, toc, section, search")
|
||||
|
||||
if final_title and final_title != title:
|
||||
result = f"Redirected to \"{final_title}\"\n\n{result}"
|
||||
|
||||
return [{"type": "text", "text": result}]
|
||||
@@ -0,0 +1,163 @@
|
||||
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
|
||||
from tools.utils import format_relative_time, ToolError, get_file_info_list, sort_file_list, get_paginated_list
|
||||
|
||||
|
||||
logger = logging.getLogger("MooseCP")
|
||||
|
||||
async def handle(args: Dict[str, Any]):
|
||||
"""
|
||||
Generates a contact sheet of images.
|
||||
Sized for optimal token usage according to config values.
|
||||
"""
|
||||
dir_path_str = args.get("path")
|
||||
page = int(args.get("page", 1))
|
||||
sort_by = args.get("sort_by", "mtime")
|
||||
|
||||
if not dir_path_str:
|
||||
raise ToolError("Missing path argument")
|
||||
|
||||
dir_path = Path(dir_path_str)
|
||||
if not dir_path.is_dir():
|
||||
raise ToolError(f"{dir_path_str} is not a directory.")
|
||||
|
||||
exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
|
||||
file_info_list = get_file_info_list(dir_path, extensions=exts)
|
||||
|
||||
if not file_info_list:
|
||||
return {"text": "No supported images found in the directory."}
|
||||
|
||||
# Sorting
|
||||
file_info_list = sort_file_list(file_info_list, sort_by)
|
||||
|
||||
sort_labels = {
|
||||
"mtime": "Date (newest first)",
|
||||
"size": "Size (largest first)",
|
||||
"name": "Name (alphabetical)"
|
||||
}
|
||||
sort_label = sort_labels.get(sort_by, "Name (alphabetical)")
|
||||
|
||||
# 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)
|
||||
paged_files = get_paginated_list(file_info_list, page, PAGE_SIZE)
|
||||
|
||||
if not paged_files:
|
||||
return {"text": f"No images found on page {page}."}
|
||||
|
||||
# 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)
|
||||
|
||||
# Font loading
|
||||
font = None
|
||||
# Try generic system font names first
|
||||
for font_name in config.SYSTEM_FONT_NAMES:
|
||||
try:
|
||||
font = ImageFont.truetype(font_name, 24)
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
# Fallback to absolute paths
|
||||
if font is None:
|
||||
for font_path in config.FALLBACK_FONT_PATHS:
|
||||
try:
|
||||
font = ImageFont.truetype(font_path, 24)
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
if font is None:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
|
||||
start_idx = (page - 1) * PAGE_SIZE
|
||||
for i, info in enumerate(paged_files):
|
||||
full_path = info["path"]
|
||||
row = i // cols
|
||||
col = i % cols
|
||||
|
||||
x = col * thumb_size
|
||||
y = row * thumb_size
|
||||
global_index = start_idx + i + 1
|
||||
|
||||
try:
|
||||
with Image.open(full_path) as img:
|
||||
# Apply EXIF orientation to fix sideways images
|
||||
img = ImageOps.exif_transpose(img)
|
||||
img.thumbnail((thumb_size, thumb_size))
|
||||
off_x = (thumb_size - img.width) // 2
|
||||
off_y = (thumb_size - img.height) // 2
|
||||
canvas.paste(img, (x + off_x, y + off_y))
|
||||
|
||||
text = str(global_index)
|
||||
if font:
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||||
else:
|
||||
text_w, text_h = len(text) * 8, 15
|
||||
|
||||
# Adjusted black box to align perfectly with the bottom of the image (y + thumb_size)
|
||||
draw.rectangle([x, y + thumb_size - text_h - 4, x + text_w + 5, y + thumb_size], fill=(0, 0, 0, 180))
|
||||
draw.text((x + 2, y + thumb_size - text_h - 7), text, fill=(255, 255, 255), font=font)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process {info['name']}: {e}")
|
||||
draw.text((x + 5, y + thumb_size // 2), "Error", fill=(255, 0, 0), font=font)
|
||||
|
||||
buf = io.BytesIO()
|
||||
# Save as JPEG to save data
|
||||
canvas.save(buf, format='JPEG', quality=config.IMAGE_QUALITY)
|
||||
img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
total_pages = (total_files + PAGE_SIZE - 1) // PAGE_SIZE
|
||||
|
||||
first_on_page = paged_files[0]
|
||||
last_on_page = paged_files[-1]
|
||||
|
||||
if sort_by == "mtime":
|
||||
first_val = f"{format_relative_time(first_on_page['mtime'])} ({datetime.datetime.fromtimestamp(first_on_page['mtime']).strftime('%Y-%m-%d')})"
|
||||
last_val = f"{format_relative_time(last_on_page['mtime'])} ({datetime.datetime.fromtimestamp(last_on_page['mtime']).strftime('%Y-%m-%d')})"
|
||||
elif sort_by == "size":
|
||||
first_val = f"{first_on_page['size']/1024:.1f}KB"
|
||||
last_val = f"{last_on_page['size']/1024:.1f}KB"
|
||||
else:
|
||||
first_val = first_on_page['name']
|
||||
last_val = last_on_page['name']
|
||||
|
||||
header_text = (
|
||||
f"Directory: {dir_path_str}\n"
|
||||
f"Page {page} of {total_pages} ({total_files} images total). Sorted by: {sort_label}\n"
|
||||
f"Range shown: {first_val} ... {last_val}"
|
||||
)
|
||||
|
||||
return [
|
||||
{"type": "text", "text": header_text},
|
||||
{"type": "image", "data": img_data, "mimeType": "image/jpeg"}
|
||||
]
|
||||
@@ -0,0 +1,245 @@
|
||||
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.
|
||||
"""
|
||||
# 1. REQUIRED ARGUMENTS
|
||||
prompt = args.get("prompt")
|
||||
if not prompt:
|
||||
raise ToolError("The 'prompt' argument is required.")
|
||||
|
||||
model_name = args.get("model_name")
|
||||
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 = 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 and a sparse payload
|
||||
label_map = {}
|
||||
label_counts = {}
|
||||
|
||||
# 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
|
||||
|
||||
# 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 = name
|
||||
|
||||
if label in label_counts:
|
||||
label_counts[label] += 1
|
||||
mapped_label = f"{label} [{label_counts[label]}]"
|
||||
else:
|
||||
label_counts[label] = 0
|
||||
mapped_label = label
|
||||
|
||||
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. 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"]:
|
||||
continue
|
||||
if key in label_map:
|
||||
payload[label_map[key]] = value
|
||||
elif key.startswith("param_"):
|
||||
try:
|
||||
idx = int(key.replace("param_", ""))
|
||||
if 0 <= idx < len(payload):
|
||||
payload[idx] = value
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
res_preset_name = args.get("resolution_preset")
|
||||
if res_preset_name:
|
||||
res_set_name = model_preset.get("Resolution Set")
|
||||
if res_set_name and res_set_name in res_cfg:
|
||||
res_set = res_cfg[res_set_name]
|
||||
if res_preset_name in res_set:
|
||||
res_val_str = res_set[res_preset_name]
|
||||
try:
|
||||
w, h = map(int, res_val_str.split('x'))
|
||||
if "Width" in label_map:
|
||||
payload[label_map["Width"]] = w
|
||||
if "Height" in label_map:
|
||||
payload[label_map["Height"]] = h
|
||||
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. 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}'.")
|
||||
|
||||
overrides = {
|
||||
"prompt": "Prompt",
|
||||
"negative_prompt": "Negative Prompt",
|
||||
"cfg_scale": "CFG Scale",
|
||||
}
|
||||
|
||||
for arg_key, label in overrides.items():
|
||||
if arg_key in args and label in label_map:
|
||||
payload[label_map[label]] = args[arg_key]
|
||||
|
||||
if "Prompt" in label_map:
|
||||
payload[label_map["Prompt"]] = prompt
|
||||
|
||||
if "Seed" in label_map:
|
||||
payload[label_map["Seed"]] = -1
|
||||
|
||||
# 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 = 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()
|
||||
|
||||
gallery_data = res_data["data"][0]["value"]
|
||||
if not gallery_data:
|
||||
raise ToolError("Server returned successfully but no image was generated.")
|
||||
|
||||
sd_img_url = gallery_data[0]["image"]["url"]
|
||||
|
||||
# Extract raw path from SD URL (e.g. http://.../file=/tmp/gradio/abc.png -> /tmp/gradio/abc.png)
|
||||
if "/file=" in sd_img_url:
|
||||
raw_path = sd_img_url.split("/file=")[1]
|
||||
else:
|
||||
# Fallback if the URL format changes
|
||||
raise ToolError(f"Could not extract file path from SD URL: {sd_img_url}")
|
||||
|
||||
# Construct proxy URL through our MCP server
|
||||
# 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 = 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')
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"Image successfully generated using model '{model_name}'.\n\nLocal path: {raw_path}\nTo show the image to the user, you can include this markdown link in your response: "
|
||||
},
|
||||
{
|
||||
"type": "image",
|
||||
"data": b64_data,
|
||||
"mimeType": "image/png"
|
||||
}
|
||||
]
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
raise ToolError(f"Server error during generation: {str(e)}")
|
||||
except Exception as e:
|
||||
raise ToolError(f"Unexpected error during generation: {str(e)}")
|
||||
@@ -0,0 +1,79 @@
|
||||
import requests
|
||||
from typing import Any, Dict, List
|
||||
import config
|
||||
from tools.utils import ToolError, load_toml
|
||||
|
||||
async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Returns information about available models or detailed guides for a specific model.
|
||||
"""
|
||||
model_name = args.get("model_name")
|
||||
|
||||
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
|
||||
catalog = []
|
||||
for name, data in models.items():
|
||||
catalog.append({
|
||||
"name": name,
|
||||
"description": data.get("description", "No description provided.")
|
||||
})
|
||||
return [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Available 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 your selected model as the 'model_name' argument."
|
||||
}
|
||||
]
|
||||
|
||||
if model_name not in models:
|
||||
# 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")
|
||||
|
||||
# Get available resolution presets for this model
|
||||
res_options = []
|
||||
if res_set_name and res_set_name in res_presets:
|
||||
res_options = list(res_presets[res_set_name].keys())
|
||||
else:
|
||||
res_options = ["Default (1024x1024)"]
|
||||
|
||||
guide = model_data.get("guide", "No detailed guide available.")
|
||||
description = model_data.get("description", "")
|
||||
|
||||
# Filter out internal config keys to show only the human-friendly presets
|
||||
presets_to_show = {k: v for k, v in model_data.items() if k not in ["description", "guide", "Resolution Set"]}
|
||||
|
||||
preset_text = "\n".join([f"- {k}: {v}" for k, v in presets_to_show.items()])
|
||||
res_text = ", ".join(res_options)
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
f"Model: {model_name}\n"
|
||||
f"Description: {description}\n\n"
|
||||
f"--- Prompting Guide ---\n{guide}\n\n"
|
||||
f"--- Active Presets ---\n{preset_text}\n\n"
|
||||
f"--- Available Resolution Presets ---\n{res_text}\n\n"
|
||||
f"Use 'resolution_preset' in generate_image to choose one of these."
|
||||
)
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,86 @@
|
||||
import re
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Set
|
||||
from tools.utils import ToolError, parse_indices
|
||||
|
||||
logger = logging.getLogger("MooseCP")
|
||||
|
||||
async def handle(args: Dict[str, Any]):
|
||||
"""
|
||||
Extracts specific sections of a file with absolute line numbers and surrounding context.
|
||||
This tool MUST be called immediately before edit_file to verify the exact line numbers
|
||||
and content of the block being replaced, preventing misalignment errors.
|
||||
"""
|
||||
path_str = args.get("path")
|
||||
pattern = args.get("pattern")
|
||||
lines_str = args.get("lines")
|
||||
context = int(args.get("context", 1))
|
||||
|
||||
if not path_str:
|
||||
raise ToolError("Missing path argument")
|
||||
|
||||
path = Path(path_str)
|
||||
if not path.is_file():
|
||||
raise ToolError(f"{path_str} is not a file.")
|
||||
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8", errors="replace") as f:
|
||||
all_lines = f.readlines()
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error reading file {path_str}: {e}")
|
||||
|
||||
total_lines = len(all_lines)
|
||||
|
||||
# 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 [{"type": "text", "text": "".join(output)}]
|
||||
|
||||
# Determine target line numbers (1-based)
|
||||
target_lines: Set[int] = set()
|
||||
|
||||
if pattern:
|
||||
try:
|
||||
regex = re.compile(pattern)
|
||||
for i, line in enumerate(all_lines):
|
||||
if regex.search(line):
|
||||
target_lines.add(i + 1)
|
||||
except re.error as e:
|
||||
raise ToolError(f"Invalid regex pattern: {e}")
|
||||
|
||||
if lines_str:
|
||||
try:
|
||||
indices = parse_indices(lines_str)
|
||||
for idx in indices:
|
||||
if 1 <= idx <= total_lines:
|
||||
target_lines.add(idx)
|
||||
except ToolError as e:
|
||||
raise e
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error parsing line indices: {e}")
|
||||
|
||||
if not target_lines:
|
||||
return [{"type": "text", "text": "No matching lines found."}]
|
||||
|
||||
# Expand targets with context
|
||||
lines_to_show: Set[int] = set()
|
||||
for line in target_lines:
|
||||
for offset in range(-context, context + 1):
|
||||
ln = line + offset
|
||||
if 1 <= ln <= total_lines:
|
||||
lines_to_show.add(ln)
|
||||
|
||||
# Sort and format output
|
||||
sorted_lines = sorted(list(lines_to_show))
|
||||
output = []
|
||||
last_line = None
|
||||
|
||||
for ln in sorted_lines:
|
||||
if last_line is not None and ln > last_line + 1:
|
||||
output.append("...\n")
|
||||
|
||||
output.append(f"{ln}: {all_lines[ln-1]}")
|
||||
last_line = ln
|
||||
|
||||
return [{"type": "text", "text": "".join(output)}]
|
||||
+25
-60
@@ -1,10 +1,9 @@
|
||||
import os
|
||||
import logging
|
||||
import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
from tools.utils import format_relative_time
|
||||
from tools.utils import format_relative_time, ToolError, get_file_info_list, sort_file_list
|
||||
|
||||
logger = logging.getLogger("MattCP")
|
||||
logger = logging.getLogger("MooseCP")
|
||||
|
||||
async def handle(args: Dict[str, Any]):
|
||||
"""
|
||||
@@ -12,65 +11,24 @@ async def handle(args: Dict[str, Any]):
|
||||
Directories are listed first, then files.
|
||||
Supports pagination and sorting.
|
||||
"""
|
||||
path = args.get("path")
|
||||
path_str = args.get("path")
|
||||
page = int(args.get("page", 1))
|
||||
page_size = int(args.get("page_size", 64))
|
||||
sort_by = args.get("sort_by", "mtime")
|
||||
|
||||
if not path:
|
||||
return {"text": "Error: Missing path argument"}
|
||||
if not os.path.isdir(path):
|
||||
return {"text": f"Error: {path} is not a directory."}
|
||||
if not path_str:
|
||||
raise ToolError("Missing path argument")
|
||||
|
||||
try:
|
||||
entries = os.listdir(path)
|
||||
except Exception as e:
|
||||
return {"text": f"Error reading directory: {e}"}
|
||||
path = Path(path_str)
|
||||
if not path.is_dir():
|
||||
raise ToolError(f"{path_str} is not a directory.")
|
||||
|
||||
all_items = []
|
||||
for entry in entries:
|
||||
full_path = os.path.join(path, entry)
|
||||
try:
|
||||
stats = os.stat(full_path)
|
||||
mtime_rel = format_relative_time(stats.st_mtime)
|
||||
|
||||
if os.path.isdir(full_path):
|
||||
try:
|
||||
count = len(os.listdir(full_path))
|
||||
except:
|
||||
count = 0
|
||||
all_items.append({
|
||||
"type": "dir",
|
||||
"name": entry,
|
||||
"info": f"{count} items",
|
||||
"mtime": mtime_rel,
|
||||
"mtime_raw": stats.st_mtime,
|
||||
"size": stats.st_size
|
||||
})
|
||||
else:
|
||||
size_kb = stats.st_size / 1024
|
||||
all_items.append({
|
||||
"type": "file",
|
||||
"name": entry,
|
||||
"info": f"{size_kb:.1f}KB",
|
||||
"mtime": mtime_rel,
|
||||
"mtime_raw": stats.st_mtime,
|
||||
"size": stats.st_size
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Could not stat {entry}: {e}")
|
||||
|
||||
# Sort Logic: Directories always come first, then apply secondary sort
|
||||
if sort_by == "mtime":
|
||||
all_items.sort(key=lambda x: (x["type"] == "file", -x["mtime_raw"]))
|
||||
elif sort_by == "size":
|
||||
all_items.sort(key=lambda x: (x["type"] == "file", -x["size"]))
|
||||
else:
|
||||
all_items.sort(key=lambda x: (x["type"] == "file", x["name"].lower()))
|
||||
all_items = get_file_info_list(path)
|
||||
all_items = sort_file_list(all_items, sort_by)
|
||||
|
||||
total_items = len(all_items)
|
||||
|
||||
# Custom Pagination: show all if <= 100, otherwise paginate by 64
|
||||
# Custom Pagination: show all if <= 100, otherwise paginate by page_size
|
||||
if total_items <= 100:
|
||||
start_idx = 0
|
||||
end_idx = total_items
|
||||
@@ -85,15 +43,22 @@ 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:
|
||||
if item["type"] == "dir":
|
||||
output.append(f"[DIR] {item['name']} | {item['info']} | {item['mtime']}")
|
||||
mtime_rel = format_relative_time(item["mtime"])
|
||||
if item["is_dir"]:
|
||||
try:
|
||||
count = len(list(item["path"].iterdir()))
|
||||
info = f"{count} items"
|
||||
except:
|
||||
info = "0 items"
|
||||
output.append(f"[DIR] {item['name']} | {info} | {mtime_rel}")
|
||||
else:
|
||||
output.append(f" {item['name']} | {item['info']} | {item['mtime']}")
|
||||
size_kb = item["size"] / 1024
|
||||
output.append(f" {item['name']} | {size_kb:.1f}KB | {mtime_rel}")
|
||||
|
||||
header = f"Listing of {path}\nPage {effective_page} of {total_pages} ({total_items} total items). Showing {start_idx+1}-{min(end_idx, total_items)}."
|
||||
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)}]
|
||||
|
||||
@@ -1,159 +0,0 @@
|
||||
import os
|
||||
import base64
|
||||
import logging
|
||||
import asyncio
|
||||
import io
|
||||
import datetime
|
||||
import time
|
||||
from typing import Any, Dict
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
logger = logging.getLogger("MattCP")
|
||||
|
||||
async def handle(args: Dict[str, Any]):
|
||||
"""
|
||||
Generates a thumbnail grid of images in a directory.
|
||||
Returns a combined text list (with global indices) and a visual contact sheet.
|
||||
"""
|
||||
dir_path = args.get("path")
|
||||
page = int(args.get("page", 1))
|
||||
page_size = int(args.get("page_size", 64))
|
||||
sort_by = args.get("sort_by", "mtime")
|
||||
|
||||
if not dir_path:
|
||||
return {"text": "Error: Missing path argument"}
|
||||
if not os.path.isdir(dir_path):
|
||||
return {"text": f"Error: {dir_path} is not a directory."}
|
||||
|
||||
exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
|
||||
file_info_list = []
|
||||
|
||||
# 1. Gather all valid images and their stats
|
||||
for f in os.listdir(dir_path):
|
||||
if f.lower().endswith(exts):
|
||||
full_path = os.path.join(dir_path, f)
|
||||
try:
|
||||
stats = os.stat(full_path)
|
||||
file_info_list.append({
|
||||
"name": f,
|
||||
"path": full_path,
|
||||
"mtime": stats.st_mtime,
|
||||
"size": stats.st_size
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Could not stat {f}: {e}")
|
||||
|
||||
if not file_info_list:
|
||||
return {"text": "No supported images found in the directory."}
|
||||
|
||||
# 2. Sort the list based on requested criteria
|
||||
if sort_by == "mtime":
|
||||
file_info_list.sort(key=lambda x: x["mtime"], reverse=True)
|
||||
elif sort_by == "size":
|
||||
file_info_list.sort(key=lambda x: x["size"], reverse=True)
|
||||
else:
|
||||
file_info_list.sort(key=lambda x: x["name"].lower())
|
||||
|
||||
# 3. Apply Pagination
|
||||
total_files = len(file_info_list)
|
||||
start_idx = (page - 1) * page_size
|
||||
end_idx = start_idx + page_size
|
||||
paged_files = file_info_list[start_idx:end_idx]
|
||||
|
||||
if not paged_files:
|
||||
return {"text": f"No images found on page {page}."}
|
||||
|
||||
# Layout constants
|
||||
thumb_size = 160
|
||||
padding = 15
|
||||
label_height = 35
|
||||
|
||||
num_files = len(paged_files)
|
||||
cols = int(num_files**0.5) if num_files > 0 else 1
|
||||
if cols == 0: cols = 1
|
||||
rows = (num_files + cols - 1) // cols
|
||||
|
||||
canvas_w = cols * (thumb_size + padding) + padding
|
||||
canvas_h = rows * (thumb_size + label_height + padding) + padding
|
||||
|
||||
# Create dark gray background
|
||||
canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30))
|
||||
draw = ImageDraw.Draw(canvas)
|
||||
|
||||
# Attempt to load a bold system font for the indices
|
||||
font = None
|
||||
font_paths = [
|
||||
"/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
|
||||
"/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
|
||||
"/usr/share/fonts/truetype/freefont/FreeSansBold.ttf",
|
||||
]
|
||||
for path in font_paths:
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
font = ImageFont.truetype(path, 20)
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
if font is None:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
file_details = []
|
||||
|
||||
# 4. Paste thumbnails and render indices
|
||||
for i, info in enumerate(paged_files):
|
||||
filename = info["name"]
|
||||
full_path = info["path"]
|
||||
row = i // cols
|
||||
col = i % cols
|
||||
|
||||
x = padding + col * (thumb_size + padding)
|
||||
y = padding + row * (thumb_size + label_height + padding)
|
||||
|
||||
# Continuous indexing across pages
|
||||
global_index = start_idx + i + 1
|
||||
|
||||
try:
|
||||
with Image.open(full_path) as img:
|
||||
width, height = img.size
|
||||
img.thumbnail((thumb_size, thumb_size))
|
||||
# Center image in its slot
|
||||
off_x = (thumb_size - img.width) // 2
|
||||
off_y = (thumb_size - img.height) // 2
|
||||
canvas.paste(img, (x + off_x, y + off_y))
|
||||
|
||||
# Format text metadata
|
||||
from tools.utils import format_relative_time # Import helper from utils
|
||||
mod_time_rel = format_relative_time(info["mtime"])
|
||||
size_kb = info["size"] / 1024
|
||||
file_details.append(f"{global_index}. {filename} | {width}x{height} | {size_kb:.1f}KB | {mod_time_rel}")
|
||||
|
||||
# Draw the index number centered under the image
|
||||
text = str(global_index)
|
||||
if font:
|
||||
bbox = draw.textbbox((0, 0), text, font=font)
|
||||
text_w = bbox[2] - bbox[0]
|
||||
else:
|
||||
text_w = len(text) * 7
|
||||
|
||||
text_x = x + (thumb_size - text_w) // 2
|
||||
draw.text((text_x, y + thumb_size + 2), text, fill=(255, 255, 255), font=font)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process {filename}: {e}")
|
||||
draw.text((x, y + thumb_size // 2), "Error", fill=(255, 0, 0), font=font)
|
||||
file_details.append(f"{global_index}. {filename} | ERROR")
|
||||
|
||||
# 5. Encode result to base64 PNG
|
||||
buf = io.BytesIO()
|
||||
canvas.save(buf, format='PNG')
|
||||
img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
total_pages = (total_files + page_size - 1) // page_size
|
||||
header_text = f"Directory: {dir_path}\nPage {page} of {total_pages} ({total_files} total images). Showing {start_idx+1}-{min(end_idx, total_files)}."
|
||||
|
||||
details_text = "\n".join(file_details)
|
||||
|
||||
return [
|
||||
{"type": "text", "text": f"{header_text}\n\nFile List:\n{details_text}"},
|
||||
{"type": "image", "data": img_data, "mimeType": "image/png"}
|
||||
]
|
||||
@@ -0,0 +1,132 @@
|
||||
import math
|
||||
import base64
|
||||
import logging
|
||||
import io
|
||||
import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Tuple
|
||||
from PIL import Image, ImageOps
|
||||
import config
|
||||
from tools.utils import format_relative_time, ToolError, get_file_info_list, sort_file_list, parse_indices
|
||||
|
||||
|
||||
logger = logging.getLogger("MooseCP")
|
||||
|
||||
def calculate_patch_dimensions(orig_w: int, orig_h: int, target_tokens: int = None):
|
||||
"""
|
||||
Calculates optimal pixel dimensions to hit a token budget of at least target_tokens.
|
||||
Validates budget based on actual rounded pixel dimensions to avoid float precision errors.
|
||||
"""
|
||||
if target_tokens is None:
|
||||
target_tokens = config.PREVIEW_TOKEN_BUDGET
|
||||
|
||||
ar = orig_w / orig_h
|
||||
|
||||
# Case A: Width is the anchor (non-padded)
|
||||
w_anchor = 1
|
||||
while True:
|
||||
w_px = w_anchor * config.PATCH_SIZE
|
||||
h_px = round(w_px / ar)
|
||||
# Calculate actual tokens based on resulting pixel dimensions
|
||||
tokens = math.ceil(w_px / config.PATCH_SIZE) * math.ceil(h_px / config.PATCH_SIZE)
|
||||
if tokens >= target_tokens:
|
||||
res_a = {"w": w_px, "h": h_px, "tokens": tokens}
|
||||
break
|
||||
w_anchor += 1
|
||||
|
||||
# Case B: Height is the anchor (non-padded)
|
||||
h_anchor = 1
|
||||
while True:
|
||||
h_px = h_anchor * config.PATCH_SIZE
|
||||
w_px = round(h_px * ar)
|
||||
# Calculate actual tokens based on resulting pixel dimensions
|
||||
tokens = math.ceil(w_px / config.PATCH_SIZE) * math.ceil(h_px / config.PATCH_SIZE)
|
||||
if tokens >= target_tokens:
|
||||
res_b = {"w": w_px, "h": h_px, "tokens": tokens}
|
||||
break
|
||||
h_anchor += 1
|
||||
|
||||
# Pick the one with the smaller token count
|
||||
best = res_a if res_a["tokens"] <= res_b["tokens"] else res_b
|
||||
return max(1, best["w"]), max(1, best["h"])
|
||||
|
||||
async def handle(args: Dict[str, Any]):
|
||||
"""
|
||||
Generates small thumbnails of images with detailed info.
|
||||
Can take a single file path, or a directory with indices/ranges from contact_sheet.
|
||||
Thumbnails are optimized for ~70 token usage.
|
||||
"""
|
||||
path_str = args.get("path")
|
||||
indices_str = args.get("indices")
|
||||
sort_by = args.get("sort_by", "mtime")
|
||||
|
||||
if not path_str:
|
||||
raise ToolError("Missing path argument")
|
||||
|
||||
path = Path(path_str)
|
||||
# Store as (path, index_label) where index_label is None if direct path
|
||||
target_files: List[Tuple[Path, Any]] = []
|
||||
|
||||
if path.is_file():
|
||||
target_files.append((path, None))
|
||||
elif path.is_dir():
|
||||
if not indices_str:
|
||||
raise ToolError("Indices are required when providing a directory path. Example: '1, 5-10'")
|
||||
|
||||
exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
|
||||
file_info_list = get_file_info_list(path, extensions=exts)
|
||||
file_info_list = sort_file_list(file_info_list, sort_by)
|
||||
|
||||
indices = parse_indices(indices_str)
|
||||
for idx in indices:
|
||||
# contact_sheet uses 1-based indexing
|
||||
if 1 <= idx <= len(file_info_list):
|
||||
target_files.append((file_info_list[idx-1]["path"], idx))
|
||||
else:
|
||||
logger.warning(f"Index {idx} out of range for directory {path_str}")
|
||||
else:
|
||||
raise ToolError(f"Path {path_str} is neither a file nor a directory.")
|
||||
|
||||
if not target_files:
|
||||
return {"text": "No valid images found to preview."}
|
||||
|
||||
results = []
|
||||
for file_path, index_label in target_files:
|
||||
try:
|
||||
stats = file_path.stat()
|
||||
with Image.open(file_path) as img:
|
||||
# Apply EXIF orientation to fix sideways images
|
||||
img = ImageOps.exif_transpose(img)
|
||||
orig_w, orig_h = img.size
|
||||
target_w, target_h = calculate_patch_dimensions(orig_w, orig_h)
|
||||
|
||||
# Use thumbnail() to preserve aspect ratio and fit within the target bounding box.
|
||||
thumb = img.convert("RGB")
|
||||
thumb.thumbnail((target_w, target_h), Image.Resampling.LANCZOS)
|
||||
|
||||
buf = io.BytesIO()
|
||||
thumb.save(buf, format='JPEG', quality=config.IMAGE_QUALITY)
|
||||
img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
# Gather info
|
||||
mtime_rel = format_relative_time(stats.st_mtime)
|
||||
mtime_abs = datetime.datetime.fromtimestamp(stats.st_mtime).strftime('%Y-%m-%d %H:%M')
|
||||
size_kb = stats.st_size / 1024
|
||||
|
||||
# Prepend index if this image was selected via index/range
|
||||
idx_prefix = f"[Index: {index_label}] " if index_label is not None else ""
|
||||
|
||||
info_text = (
|
||||
f"{idx_prefix}File: {file_path.name}\n"
|
||||
f"Resolution: {orig_w}x{orig_h}\n"
|
||||
f"Size: {size_kb:.1f}KB\n"
|
||||
f"Modified: {mtime_rel} ({mtime_abs})"
|
||||
)
|
||||
|
||||
results.append({"type": "text", "text": info_text})
|
||||
results.append({"type": "image", "data": img_data, "mimeType": "image/jpeg"})
|
||||
|
||||
except Exception as e:
|
||||
results.append({"type": "text", "text": f"Error processing {file_path.name}: {e}"})
|
||||
|
||||
return results
|
||||
+13
-12
@@ -1,29 +1,30 @@
|
||||
import os
|
||||
import base64
|
||||
import mimetypes
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
from tools.utils import ToolError
|
||||
|
||||
logger = logging.getLogger("MattCP")
|
||||
logger = logging.getLogger("MooseCP")
|
||||
|
||||
async def handle(args: Dict[str, Any]):
|
||||
"""
|
||||
Reads an image file from disk and returns it as a base64 encoded string
|
||||
within an MCP ImageContent object.
|
||||
"""
|
||||
path = args.get("path")
|
||||
if not path:
|
||||
return {"text": "Error: Missing path argument"}
|
||||
path_str = args.get("path")
|
||||
if not path_str:
|
||||
raise ToolError("Missing path argument")
|
||||
|
||||
path = Path(path_str)
|
||||
|
||||
# Validate that the file has an image-like MIME type
|
||||
mime_type, _ = mimetypes.guess_type(path)
|
||||
mime_type, _ = mimetypes.guess_type(path_str)
|
||||
if not mime_type or not mime_type.startswith("image/"):
|
||||
return {"text": f"Error: File {path} is not a supported image format."}
|
||||
raise ToolError(f"File {path_str} is not a supported image format.")
|
||||
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
# Read raw bytes and encode to base64 for transport
|
||||
data = base64.b64encode(f.read()).decode("utf-8")
|
||||
return {"type": "image", "data": data, "mimeType": mime_type}
|
||||
data = base64.b64encode(path.read_bytes()).decode("utf-8")
|
||||
return [{"type": "image", "data": data, "mimeType": mime_type}]
|
||||
except Exception as e:
|
||||
return {"text": f"Error reading image: {str(e)}"}
|
||||
raise ToolError(f"Error reading image: {str(e)}")
|
||||
|
||||
+14
-9
@@ -1,32 +1,37 @@
|
||||
import os
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
from PIL import Image
|
||||
from tools.utils import ToolError
|
||||
|
||||
logger = logging.getLogger("MattCP")
|
||||
logger = logging.getLogger("MooseCP")
|
||||
|
||||
async def handle(args: Dict[str, Any]):
|
||||
"""
|
||||
Extracts text-based metadata chunks (tEXt, zTXt, iTXt) from a PNG file.
|
||||
These often contain AI generation prompts and parameters.
|
||||
"""
|
||||
path = args.get("path")
|
||||
if not path:
|
||||
return {"text": "Error: Missing path argument"}
|
||||
path_str = args.get("path")
|
||||
if not path_str:
|
||||
raise ToolError("Missing path argument")
|
||||
|
||||
path = Path(path_str)
|
||||
|
||||
try:
|
||||
with Image.open(path) as img:
|
||||
# Metadata extraction is specific to PNG format in this implementation
|
||||
if img.format != 'PNG':
|
||||
return {"text": "Error: File is not a PNG image."}
|
||||
raise ToolError("File is not a PNG image.")
|
||||
|
||||
# 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:
|
||||
return {"text": f"Error reading PNG metadata: {str(e)}"}
|
||||
raise ToolError(f"Error reading PNG metadata: {str(e)}")
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
import csv
|
||||
import os
|
||||
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
|
||||
|
||||
# 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.")
|
||||
|
||||
try:
|
||||
with open(tag_file_path, mode='r', encoding='utf-8') as f:
|
||||
reader = csv.reader(f)
|
||||
for row in reader:
|
||||
if not row or len(row) < 3:
|
||||
continue
|
||||
|
||||
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
|
||||
|
||||
return tags
|
||||
|
||||
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:
|
||||
output_lines.append("") # spacer
|
||||
|
||||
# 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 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
|
||||
|
||||
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)}]
|
||||
+108
@@ -1,5 +1,21 @@
|
||||
import time
|
||||
import datetime
|
||||
import tomllib
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
|
||||
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."""
|
||||
@@ -32,3 +48,95 @@ def format_relative_time(timestamp: float) -> str:
|
||||
|
||||
dt = datetime.datetime.fromtimestamp(timestamp)
|
||||
return dt.strftime('%Y-%m-%d')
|
||||
|
||||
|
||||
def get_file_info_list(path: Path, extensions: Optional[tuple] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Gathers basic information about files in a directory.
|
||||
Optional extensions filter (e.g. ('.png', '.jpg')).
|
||||
"""
|
||||
items = []
|
||||
try:
|
||||
for entry in path.iterdir():
|
||||
if extensions and not entry.name.lower().endswith(extensions):
|
||||
continue
|
||||
|
||||
stats = entry.stat()
|
||||
items.append({
|
||||
"name": entry.name,
|
||||
"path": entry,
|
||||
"mtime": stats.st_mtime,
|
||||
"size": stats.st_size,
|
||||
"is_dir": entry.is_dir()
|
||||
})
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error accessing directory {path}: {e}")
|
||||
|
||||
return items
|
||||
|
||||
def sort_file_list(items: List[Dict[str, Any]], sort_by: str) -> List[Dict[str, Any]]:
|
||||
"""Sorts items based on the provided key. Directories always come first."""
|
||||
if sort_by == "mtime":
|
||||
# Newest first
|
||||
items.sort(key=lambda x: (not x["is_dir"], -x["mtime"]))
|
||||
elif sort_by == "size":
|
||||
# Largest first
|
||||
items.sort(key=lambda x: (not x["is_dir"], -x["size"]))
|
||||
else:
|
||||
# Alphabetical
|
||||
items.sort(key=lambda x: (not x["is_dir"], x["name"].lower()))
|
||||
return items
|
||||
|
||||
def get_paginated_list(items: List[Any], page: int, page_size: int) -> List[Any]:
|
||||
"""Returns a slice of the list based on page and page_size."""
|
||||
start_idx = (page - 1) * page_size
|
||||
end_idx = start_idx + page_size
|
||||
return items[start_idx:end_idx]
|
||||
|
||||
def parse_indices(indices_str: str) -> List[int]:
|
||||
"""
|
||||
Parses a string like '11, 31-33, 44' into a list of 1-based indices.
|
||||
Preserves the order specified in the string.
|
||||
"""
|
||||
indices = []
|
||||
parts = [p.strip() for p in indices_str.split(',')]
|
||||
for part in parts:
|
||||
if not part:
|
||||
continue
|
||||
if '-' in part:
|
||||
try:
|
||||
start, end = map(int, part.split('-'))
|
||||
step = 1 if start <= end else -1
|
||||
for i in range(start, end + step, step):
|
||||
indices.append(i)
|
||||
except ValueError:
|
||||
raise ToolError(f"Invalid range format: {part}")
|
||||
else:
|
||||
try:
|
||||
indices.append(int(part))
|
||||
except ValueError:
|
||||
raise ToolError(f"Invalid index format: {part}")
|
||||
return indices
|
||||
|
||||
def format_count(count_str: str) -> str:
|
||||
"""Formats large numbers into a human-readable string (e.g., 1.2M, 218k)."""
|
||||
try:
|
||||
count = int(count_str)
|
||||
if count >= 1_000_000:
|
||||
return f"{count / 1_000_000:.1f}M".replace(".0", "")
|
||||
if count >= 1_000:
|
||||
return f"{count / 1_000:.0f}k"
|
||||
return str(count)
|
||||
except (ValueError, TypeError):
|
||||
return "0"
|
||||
|
||||
def get_type_suffix(type_val: str) -> str:
|
||||
"""Returns the human-readable suffix for the tag type."""
|
||||
mapping = {
|
||||
"0": "", # General
|
||||
"1": " [Artist]",
|
||||
"3": " [Copyright]",
|
||||
"4": " [Character]",
|
||||
"5": " [Meta]"
|
||||
}
|
||||
return mapping.get(str(type_val), f" [Type {type_val}]")
|
||||
|
||||
Reference in New Issue
Block a user