Code review fixes: wardrobe migration, response validation, path traversal guard, deduplication

- Migrate 11 character JSONs from old wardrobe keys to _BODY_GROUP_KEYS format
- Add is_favourite/is_nsfw columns to Preset model
- Add HTTP response validation and timeouts to ComfyUI client
- Add path traversal protection on replace cover route
- Deduplicate services/mcp.py (4 functions → 2 generic + 2 wrappers)
- Extract apply_library_filters() and clean_html_text() shared helpers
- Add named constants for 17 ComfyUI workflow node IDs
- Fix bare except clauses in services/llm.py
- Fix tags schema in ensure_default_outfit() (list → dict)
- Convert f-string logging to lazy % formatting
- Add 5-minute polling timeout to frontend waitForJob()
- Improve migration error handling (non-duplicate errors log at WARNING)
- Update CLAUDE.md to reflect all changes

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Aodhan Collins
2026-03-22 00:31:27 +00:00
parent 55ff58aba6
commit 29a6723b25
37 changed files with 464 additions and 539 deletions

View File

@@ -30,7 +30,7 @@ services/
generation.py # Shared generation logic (generate_from_preset)
routes/
__init__.py # register_routes(app) — imports and calls all route modules
shared.py # Factory functions for common routes (favourite, upload, clone, save_json, etc.)
shared.py # Factory functions for common routes + apply_library_filters() helper
characters.py # Character CRUD + generation + outfit management
outfits.py # Outfit routes
actions.py # Action routes
@@ -115,27 +115,29 @@ All category models (except Settings and Checkpoint) share this pattern:
## ComfyUI Workflow Node Map
The workflow (`comfy_workflow.json`) uses string node IDs. These are the critical nodes:
The workflow (`comfy_workflow.json`) uses string node IDs. Named constants are defined in `services/workflow.py`:
| Node | Role |
|------|------|
| `3` | Main KSampler |
| `4` | Checkpoint loader |
| `5` | Empty latent (width/height) |
| `6` | Positive prompt — contains `{{POSITIVE_PROMPT}}` placeholder |
| `7` | Negative prompt |
| `8` | VAE decode |
| `9` | Save image |
| `11` | Face ADetailer |
| `13` | Hand ADetailer |
| `14` | Face detailer prompt — contains `{{FACE_PROMPT}}` placeholder |
| `15` | Hand detailer prompt — contains `{{HAND_PROMPT}}` placeholder |
| `16` | Character LoRA (or Look LoRA when a Look is active) |
| `17` | Outfit LoRA |
| `18` | Action LoRA |
| `19` | Style / Detailer / Scene LoRA (priority: style > detailer > scene) |
| Constant | Node | Role |
|----------|------|------|
| `NODE_KSAMPLER` | `3` | Main KSampler |
| `NODE_CHECKPOINT` | `4` | Checkpoint loader |
| `NODE_LATENT` | `5` | Empty latent (width/height) |
| `NODE_POSITIVE` | `6` | Positive prompt — contains `{{POSITIVE_PROMPT}}` placeholder |
| `NODE_NEGATIVE` | `7` | Negative prompt |
| `NODE_VAE_DECODE` | `8` | VAE decode |
| `NODE_SAVE` | `9` | Save image |
| `NODE_FACE_DETAILER` | `11` | Face ADetailer |
| `NODE_HAND_DETAILER` | `13` | Hand ADetailer |
| `NODE_FACE_PROMPT` | `14` | Face detailer prompt — contains `{{FACE_PROMPT}}` placeholder |
| `NODE_HAND_PROMPT` | `15` | Hand detailer prompt — contains `{{HAND_PROMPT}}` placeholder |
| `NODE_LORA_CHAR` | `16` | Character LoRA (or Look LoRA when a Look is active) |
| `NODE_LORA_OUTFIT` | `17` | Outfit LoRA |
| `NODE_LORA_ACTION` | `18` | Action LoRA |
| `NODE_LORA_STYLE` | `19` | Style / Detailer / Scene LoRA (priority: style > detailer > scene) |
| `NODE_LORA_CHAR_B` | `20` | Character LoRA B (second character) |
| `NODE_VAE_LOADER` | `21` | VAE loader |
LoRA nodes chain: `4 → 16 → 17 → 18 → 19`. Unused LoRA nodes are bypassed by pointing `model_source`/`clip_source` directly to the prior node. All model/clip consumers (nodes 3, 6, 7, 11, 13, 14, 15) are wired to the final `model_source`/`clip_source` at the end of `_prepare_workflow`.
LoRA nodes chain: `4 → 16 → 17 → 18 → 19`. Unused LoRA nodes are bypassed by pointing `model_source`/`clip_source` directly to the prior node. All model/clip consumers (nodes 3, 6, 7, 11, 13, 14, 15) are wired to the final `model_source`/`clip_source` at the end of `_prepare_workflow`. Always use the named constants instead of string literals when referencing node IDs.
---
@@ -143,12 +145,13 @@ LoRA nodes chain: `4 → 16 → 17 → 18 → 19`. Unused LoRA nodes are bypasse
### `utils.py` — Constants and Pure Helpers
- **`_IDENTITY_KEYS` / `_WARDROBE_KEYS`** — Lists of canonical field names for the `identity` and `wardrobe` sections. Used by `_ensure_character_fields()`.
- **`_BODY_GROUP_KEYS`** — Canonical list of field names shared by both `identity` and `wardrobe` sections: `['base', 'head', 'upper_body', 'lower_body', 'hands', 'feet', 'additional']`. Used by `build_prompt()`, `_ensure_character_fields()`, and `_resolve_preset_fields()`.
- **`ALLOWED_EXTENSIONS`** — Permitted upload file extensions.
- **`_LORA_DEFAULTS`** — Default LoRA directory paths per category.
- **`parse_orientation(orientation_str)`** — Converts orientation codes (`1F`, `2F`, `1M1F`, etc.) into Danbooru tags.
- **`_resolve_lora_weight(lora_data)`** — Extracts and validates LoRA weight from a lora data dict.
- **`allowed_file(filename)`** — Checks file extension against `ALLOWED_EXTENSIONS`.
- **`clean_html_text(html_raw)`** — Strips HTML tags, scripts, styles, and images from raw HTML, returning plain text. Used by bulk_create routes.
### `services/prompts.py` — Prompt Building
@@ -181,10 +184,11 @@ Two independent queues with separate worker threads:
### `services/comfyui.py` — ComfyUI HTTP Client
- **`queue_prompt(prompt_workflow, client_id)`** — POSTs workflow to ComfyUI's `/prompt` endpoint.
- **`get_history(prompt_id)`** — Polls ComfyUI for job completion.
- **`get_image(filename, subfolder, folder_type)`** — Retrieves generated image bytes.
- **`_ensure_checkpoint_loaded(checkpoint_path)`** — Forces ComfyUI to load a specific checkpoint.
- **`queue_prompt(prompt_workflow, client_id)`** — POSTs workflow to ComfyUI's `/prompt` endpoint. Validates HTTP response status before parsing JSON; raises `RuntimeError` on non-OK responses. Timeout: 30s.
- **`get_history(prompt_id)`** — Polls ComfyUI for job completion. Timeout: 10s.
- **`get_image(filename, subfolder, folder_type)`** — Retrieves generated image bytes. Timeout: 30s.
- **`get_loaded_checkpoint()`** — Returns the checkpoint path currently loaded in ComfyUI by inspecting the most recent job in `/history`.
- **`_ensure_checkpoint_loaded(checkpoint_path)`** — Forces ComfyUI to unload all models if the desired checkpoint doesn't match what's currently loaded.
### `services/llm.py` — LLM Integration
@@ -211,8 +215,10 @@ Two independent queues with separate worker threads:
### `services/mcp.py` — MCP/Docker Lifecycle
- **`ensure_mcp_server_running()`** — Ensures the danbooru-mcp Docker container is running.
- **`ensure_character_mcp_server_running()`** — Ensures the character-mcp Docker container is running.
- **`_ensure_repo(compose_dir, repo_url, name)`** — Generic helper: clones a git repo if the directory doesn't exist.
- **`_ensure_server_running(compose_dir, repo_url, container_name, name)`** — Generic helper: ensures a Docker Compose service is running (clones repo if needed, starts container if not running).
- **`ensure_mcp_server_running()`** — Ensures the danbooru-mcp Docker container is running (thin wrapper around `_ensure_server_running`).
- **`ensure_character_mcp_server_running()`** — Ensures the character-mcp Docker container is running (thin wrapper around `_ensure_server_running`).
### Route-local Helpers
@@ -239,7 +245,7 @@ Some helpers are defined inside a route module's `register_routes()` since they'
"identity": { "base_specs": "", "hair": "", "eyes": "", "hands": "", "arms": "", "torso": "", "pelvis": "", "legs": "", "feet": "", "extra": "" },
"defaults": { "expression": "", "pose": "", "scene": "" },
"wardrobe": {
"default": { "full_body": "", "headwear": "", "top": "", "bottom": "", "legwear": "", "footwear": "", "hands": "", "gloves": "", "accessories": "" }
"default": { "base": "", "head": "", "upper_body": "", "lower_body": "", "hands": "", "feet": "", "additional": "" }
},
"styles": { "aesthetic": "", "primary_color": "", "secondary_color": "", "tertiary_color": "" },
"lora": { "lora_name": "Illustrious/Looks/tifa.safetensors", "lora_weight": 0.8, "lora_triggers": "" },
@@ -254,7 +260,7 @@ Some helpers are defined inside a route module's `register_routes()` since they'
{
"outfit_id": "french_maid_01",
"outfit_name": "French Maid",
"wardrobe": { "full_body": "", "headwear": "", "top": "", "bottom": "", "legwear": "", "footwear": "", "hands": "", "accessories": "" },
"wardrobe": { "base": "", "head": "", "upper_body": "", "lower_body": "", "hands": "", "feet": "", "additional": "" },
"lora": { "lora_name": "Illustrious/Clothing/maid.safetensors", "lora_weight": 0.8, "lora_triggers": "" },
"tags": { "outfit_type": "Uniform", "nsfw": false }
}
@@ -439,7 +445,7 @@ Image retrieval is handled server-side by the `_make_finalize()` callback; there
- `static/js/library-toolbar.js` — Library page toolbar (batch generate, clear covers, missing items)
- Context processors inject `all_checkpoints`, `default_checkpoint_path`, and `COMFYUI_WS_URL` into every template. The `random_gen_image(category, slug)` template global returns a random image path from `static/uploads/<category>/<slug>/` for use as a fallback cover when `image_path` is not set.
- **No `{% block head %}` exists** in layout.html — do not try to use it.
- Generation is async: JS submits the form via AJAX (`X-Requested-With: XMLHttpRequest`), receives a `{"job_id": ...}` response, then polls `/api/queue/<job_id>/status` every ~1.5 seconds until `status == "done"`. The server-side worker handles all ComfyUI polling and image saving via the `_make_finalize()` callback. There are no client-facing finalize HTTP routes.
- Generation is async: JS submits the form via AJAX (`X-Requested-With: XMLHttpRequest`), receives a `{"job_id": ...}` response, then polls `/api/queue/<job_id>/status` every ~1.5 seconds until `status == "done"` or the 5-minute timeout is reached. The server-side worker handles all ComfyUI polling and image saving via the `_make_finalize()` callback. There are no client-facing finalize HTTP routes.
- **Batch generation** (library pages): Uses a two-phase pattern:
1. **Queue phase**: All jobs are submitted upfront via sequential fetch calls, collecting job IDs
2. **Poll phase**: All jobs are polled concurrently via `Promise.all()`, updating UI as each completes
@@ -501,6 +507,7 @@ All library index pages support query params:
- `?favourite=on` — show only favourites
- `?nsfw=sfw|nsfw|all` — filter by NSFW status
- Results are ordered by `is_favourite DESC, name ASC` (favourites sort first).
- Filter logic is shared via `apply_library_filters(query, model_class)` in `routes/shared.py`, which returns `(items, fav, nsfw)`.
### Gallery Image Sidecar Files
@@ -609,3 +616,5 @@ Volumes mounted into the app container:
- **`_make_finalize` action semantics**: Pass `action=None` when the route should always update the DB cover (e.g. batch generate, checkpoint generate). Pass `action=request.form.get('action')` for routes that support both "preview" (no DB update) and "replace" (update DB). The factory skips the DB write when `action` is truthy and not `"replace"`.
- **LLM queue runs without request context**: `_enqueue_task()` callbacks execute in a background thread with only `app.app_context()`. Do not access `flask.request`, `flask.session`, or other request-scoped objects inside `task_fn`. Use `has_request_context()` guard if code is shared between HTTP handlers and background tasks.
- **Tags are metadata only**: Tags (`data['tags']`) are never injected into generation prompts. They are purely for UI filtering and search. The old pattern of `parts.extend(data.get('tags', []))` in prompt building has been removed.
- **Path traversal guard on replace cover**: The replace cover route in `routes/shared.py` validates `preview_path` using `os.path.realpath()` + `startswith()` to prevent path traversal attacks.
- **Logging uses lazy % formatting**: All logger calls use `logger.info("msg %s", var)` style, not f-strings. This avoids formatting the string when the log level is disabled.