diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..b7c9c02 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,51 @@ +name: CI + +on: + push: + pull_request: + +jobs: + test-and-lint: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Configure Git safe directory + run: git config --global safe.directory "*" + + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + python -m pip install pytest black ruff mypy + + - name: Block Gradio .update regression + run: python tools/check_no_gradio_update.py + + - name: Run tests (Gradio 3 stub) + run: python -m pytest tests/ -q + + - name: Run tests (Gradio 4 stub) + run: python -m pytest tests/ -q --gradio-version=4 + + - name: Ruff + run: python -m ruff check scripts/ranbooru.py ranboorux tests tools install.py + + - name: Black + run: python -m black --check scripts/ranbooru.py ranboorux tests tools install.py + + - name: Mypy + run: python -m mypy ranboorux --warn-return-any --warn-unused-ignores diff --git a/.gitignore b/.gitignore index f92d164..cbc5fac 100644 --- a/.gitignore +++ b/.gitignore @@ -2,10 +2,13 @@ .vscode/ .idea/ *.code-workspace +.project +.settings/ # OS clutter .DS_Store Thumbs.db +desktop.ini # Python cache/build/test artifacts __pycache__/ @@ -13,13 +16,13 @@ __pycache__/ *.pyo *.pyd .pytest_cache/ +.pytest_cache_local/ .mypy_cache/ .ruff_cache/ .hypothesis/ .pyre/ .tox/ .nox/ -.pytest_cache_local/ .cache/ .eggs/ *.egg-info/ @@ -31,13 +34,16 @@ dist/ pip-wheel-metadata/ __pypackages__/ +# SQLite databases & temp test artifacts +*.sqlite +*.db +temp_test_cache.sqlite + # Local env/secrets .env .env.* *.env .python-version -.gelbooru_credentials.json -gelbooru_creds.json # Virtual environments .venv/ @@ -45,15 +51,12 @@ venv/ env/ ENV/ -# Logs/runtime/databases +# Logs/runtime *.log logs/ *.pid -*.sqlite -*.sqlite3 -temp_test_cache.sqlite -# Temporary files & junk +# Temporary files *.tmp *.bak *.swp @@ -65,24 +68,28 @@ tmpclaude-* tmp_compile_err.txt temp_*.py -# Local tooling / personal notes / standalone scripts +# Agent / Tooling runtime clutter +.omo/ .roomodes .rooroo/ -docs/ +.agent/ +.agents/ +.gemini/ +.claude/ cleancode.md ranbooru.before_revert.py block.txt line241.txt tools/debug_lines.py -adetailer/ -scripts/comments.py -.antigravityignore -.brain/ + +docs/ # Extension user data user/ -# Bundled catalog assets are tracked +# Ignore nested extension directory +adetailer/ + +# Track bundled catalog assets !data/catalogs/ !data/catalogs/** -.github/workflows/ci.yml diff --git a/README.md b/README.md index 1fa7105..ddb0afd 100644 --- a/README.md +++ b/README.md @@ -4,24 +4,26 @@ ![RanbooruX logo](pics/ranbooru.png) -RanbooruX is a fork of Ranbooru focused on **Forge Neo**. +RanbooruX is a fork of Ranbooru built **exclusively for Forge Neo**, featuring native support for **ADetailer Neo**. -It fetches booru tags and source images, builds prompts, and supports a two-stage generation flow with optional Img2Img, ControlNet handoff, and ADetailer postprocessing. +It fetches booru tags and source images, builds prompts, and supports a two-stage generation flow with Img2Img, ControlNet handoff, and ADetailer / ADetailer Neo postprocessing. -## Platform support +![UI screenshot](pics/image.jpg) -> [!IMPORTANT] -> **Project Owner Testing Disclaimer**: This project is strictly developed and tested **only using Forge Neo**. Other WebUI distributions (including original SD WebUI / Automatic1111 and original SD WebUI Forge) are **not tested** by the repository owner. Support for other environments is strictly best-effort. +## Features & Exclusive Fork Capabilities -## Why this fork? +RanbooruX delivers massive architectural and feature upgrades over original Ranbooru: -- Fix brittle Img2Img/ControlNet interactions and make them reliable on **Forge Neo**. -- Split the old “remove bad tags” into clear, no‑surprise filters. -- Make installs easy with `requirements.txt` and a bundled ControlNet helper. -- Add favorites, file‑driven prompts, logging, and sensible caching. -- ![UI screenshot](pics/image.png) +- **Forge Neo & ADetailer Neo Native Support**: Built exclusively for Forge Neo, with full support for ADetailer Neo and standard ADetailer in two-pass Img2Img workflows. +- **Anima (2B DiT) Support**: Native auto-detection of Anima models with automatic flow-matching scheduler tuning, prompt quality prefixes, and working basic Img2Img & ControlNet LLLite support. +- **Danbooru Tag Catalog System**: Bundled tag catalog (`data/catalogs/danbooru_tags.csv`) providing alias normalization, category-aware filtering, custom CSV import, and hair/eye color preservation. +- **Safer Two-Pass Img2Img & Guarded Postprocessing**: Preview guard suppresses initial-pass flashes until final img2img outputs are rendered; guarded script runner prevents script collisions. +- **Rich Booru & Tag Removal Filters**: Multi-booru search (`aibooru`, `danbooru`, `e621`, `gelbooru`, `konachan`, `rule34`, `safebooru`, `xbooru`, `yande.re`) with fine-grained removal toggles (artist, character, series, clothing, commentary, furry, headwear, `*_girl` suffix cleanup). +- **LoRAnado Random LoRA Injection**: Automatic detection and control surfaces for PonyXL & Anima-compatible LoRAs with blacklist support. +- **Modular Codebase & Quality Tooling**: Refactored from a monolithic script into a clean `ranboorux/` module with unit tests (`pytest`), strict type checking (`mypy`), linting (`ruff`), and formatting (`black`). +- **User Conveniences**: Favorites management, file-driven tag sources, prompt/source logging, and sensible caching. ## Installation @@ -36,127 +38,227 @@ It fetches booru tags and source images, builds prompts, and supports a two-stag ### Method 2: Manual Installation -1. Copy or clone this repository to your **Forge Neo** extensions directory: +1. Copy or clone this repository to your WebUI extensions directory: - `extensions/sd-webui-ranbooruX` -2. Start or restart **Forge Neo**. +2. Start or restart WebUI. 3. `install.py` installs extension dependencies from `requirements.txt`. 4. Open the **RanbooruX** panel. -Optional environment overrides for ControlNet detection: +
+Environment Configuration Overrides -- `SD_FORGE_CONTROLNET_PATH` -- `RANBOORUX_CN_PATH` +If your Forge Neo ControlNet extension is located in a custom or non-standard directory outside of `extensions/sd_forge_controlnet` or `extensions-builtin/sd_forge_controlnet`, RanbooruX supports optional environment variables to override the ControlNet detection path: -## Quick start +* `SD_FORGE_CONTROLNET_PATH`: Primary override pointing to the root directory of `sd_forge_controlnet` (containing `lib_controlnet/external_code.py`). +* `RANBOORUX_CN_PATH`: Secondary fallback override path for RanbooruX ControlNet asset resolution. -1. Select a booru source. -2. Enter tags and generate. -3. Optional: enable `Use Image for Img2Img`. -4. Optional: enable `Use Image for ControlNet (Unit 0)`. -5. Optional: enable `Enable RanbooruX ADetailer support`. +#### How to Configure -## Key features +##### Option 1: In WebUI Startup Scripts (Recommended) +Add the environment variable directly to your WebUI launcher so it persists across restarts: -- Booru sources: `aibooru`, `danbooru`, `e621`, `gelbooru`, `gelbooru-compatible`, `konachan`, `rule34`, `safebooru`, `xbooru`, `yande.re` -- Fine-grained removal filters (artist, character, series, clothing, text/commentary, furry, headwear, `*_girl`, subject constraints, preserve hair/eye colors, and more) -- `Quick Strip` one-click removal preset (instantly activates all major removal filters for aggressive prompt cleanup) -- Danbooru tag catalog normalization/filtering (enabled by default, toggleable) -- Img2Img and ControlNet handoff flow -- Optional manual ADetailer pass after Img2Img -- LoRAnado random LoRA injection with PonyXL & Anima compatibility controls (legacy feature) -- Platform diagnostics panel for runtime visibility -- Caching, file-driven tag sources, favorites, and prompt/source logging +* **Windows (`webui-user.bat`)**: + ```cmd + set SD_FORGE_CONTROLNET_PATH=C:\path\to\sd_forge_controlnet + ``` +* **Linux / macOS (`webui-user.sh`)**: + ```bash + export SD_FORGE_CONTROLNET_PATH="/path/to/sd_forge_controlnet" + ``` -## Danbooru Tag Catalog +##### Option 2: Terminal / Shell -RanbooruX includes a bundled catalog used by the redesigned tag-catalog pipeline. +* **Windows PowerShell**: + ```powershell + $env:SD_FORGE_CONTROLNET_PATH="C:\path\to\sd_forge_controlnet" + ``` +* **Windows Command Prompt (`cmd.exe`)**: + ```cmd + set SD_FORGE_CONTROLNET_PATH=C:\path\to\sd_forge_controlnet + ``` +* **Linux / macOS (`bash` / `zsh`)**: + ```bash + export SD_FORGE_CONTROLNET_PATH="/path/to/sd_forge_controlnet" + ``` + +
+ +## Quick Start & Usage Workflow + +RanbooruX integrates prompt fetching, tag catalog processing, two-pass generation, ControlNet handoff, and postprocessing into a streamlined workflow: + +1. **Select Source & Query Tags**: Choose a booru source (`danbooru`, `gelbooru`, `e621`, etc.), enter your desired search tags, and specify post limits. +2. **Apply Tag Catalog & Removal Filters**: Keep `Use Danbooru Tag Catalog` enabled (default ON) for alias normalization and category-aware filtering. Configure tag removal toggles to strip unwanted artist/character/clothing metadata. +3. **Configure Image Handoff (Optional)**: + - Check `Use Image for Img2Img` to run an initial pass followed by an Img2Img pass with automatic denoising caps. + - Check `Use Image for ControlNet (Unit 0)` to automatically pass the fetched booru image into Forge Neo's ControlNet Unit 0 slot. +4. **Enable Postprocessing (Optional)**: Check `Enable RanbooruX ADetailer support` to automatically run a guarded manual ADetailer or ADetailer Neo pass on the final outputs. +5. **Generate**: Click **Generate** — RanbooruX fetches posts, processes prompts, and executes the multi-pass pipeline automatically. + +## Tag Filters & Catalog Processing + +RanbooruX provides powerful tag filtering and catalog normalization to keep prompts clean and coherent. + +![Tag Removal Filters](pics/filters.jpg) + +### Danbooru Tag Catalog + +RanbooruX includes a bundled catalog used by the tag-catalog pipeline. - Bundled file: `data/catalogs/danbooru_tags.csv` - Catalog mode toggle: `Use Danbooru Tag Catalog` (default ON) - Source selection: `Bundled` or `Custom file` -With catalog mode enabled (default), the catalog pipeline adds: +With catalog mode enabled (default), the pipeline provides: +- Alias normalization (maps variant tags to canonical Danbooru tags) +- Category-aware tag filtering (artist, character, series, meta, general) +- Smart hair & eye color preservation +- Textual & commentary tag cleanup -- alias normalization -- category-aware filtering -- better hair/eye preservation behavior -- textual/meta tag cleanup backed by catalog categories -- diagnostics panel for kept/dropped/unknown tag insight +### Custom Catalog Files -When the toggle is disabled, RanbooruX still uses the bundled catalog path (catalog-only mode; no legacy filter engine). - -### Custom catalog files - -Custom CSV catalogs are supported and imported into `user/catalogs/`. +Custom CSV catalogs can be imported into `user/catalogs/` via the UI. Accepted formats: - - Header-based CSV (`tag,category,count,alias`) - Headerless 4-column CSV (`tag,category,count,alias`) -Validation/import controls: +Validation and management buttons (`Validate CSV`, `Import Custom Catalog`, `Reload Catalog`) are provided in the UI. -- `Validate CSV` -- `Import Custom Catalog` -- `Reload Catalog` +## Two-Pass Img2Img + ADetailer / ADetailer Neo Pipeline -Implementation details and format notes are documented in: +For Img2Img workflows, RanbooruX executes a coordinated multi-stage process: -- `data/catalogs/README.txt` -- `ranboorux/catalog.py` +1. **Initial Pass & Preview Guard**: Generates the base image while suppressing intermediate preview flashes until final images are rendered. +2. **Img2Img Pass**: Automatically applies tuned denoising caps to refine details without breaking composition. +3. **ControlNet Handoff**: When `Use Image for ControlNet (Unit 0)` is enabled, the fetched booru reference image is automatically assigned to Unit 0 in Forge Neo's ControlNet runner. +4. **ADetailer / ADetailer Neo Pass**: Runs a guarded postprocessing pass on the final images, auto-detecting both standard ADetailer and ADetailer Neo scripts. -### Bundled catalog provenance and licensing notes +## Anima Model Support -`data/catalogs/README.txt` includes provenance/licensing context for the bundled `danbooru_tags.csv`, plus references used for the research notes. +RanbooruX natively supports **Anima** (a 2B parameter DiT model by CircleStone Labs + Comfy Org built on NVIDIA Cosmos-Predict2) in Forge Neo with basic Img2Img support fully working. -## LoRAnado (PonyXL & Anima detection) +### Anima ControlNet Support -> [!NOTE] -> LoRAnado is a legacy feature inherited from original Ranbooru and is not extensively tested by the repository owner. +RanbooruX supports **basic ControlNet Img2Img & conditioning handoff** for Anima models. -LoRAnado includes detection and control surfaces to reduce incompatible LoRA picks in PonyXL and Anima workflows. +Anima uses a 2B Diffusion Transformer (DiT) architecture, which requires specialized **ControlNet-LLLite** models rather than standard SD/SDXL ControlNets: -Controls: +- **Available LLLite Models**: `anima-lllite-lineart-1` (line art / pose guidance), `anima-lllite-depth-1` (depth estimation guidance), `anima-lllite-inpainting-v2` (targeted inpainting). +- **How to Use**: + 1. Open Forge Neo's **ControlNet** panel (Unit 0 tab). + 2. Select an Anima LLLite model (`anima-lllite-lineart-1` or `anima-lllite-depth-1`) and matching preprocessor (`anime_lineart` or `depth`). + 3. In RanbooruX, check **`Use Image for ControlNet (Unit 0)`**. + 4. Click **Generate** — RanbooruX automatically passes the fetched booru image to Unit 0. +- **Scope & Limitations**: RanbooruX handles standard ControlNet LLLite image handoff into Unit 0. Anima Edit (Cosmos-Reference) is not supported. -- `Auto-detect PonyXL/Anima-compatible LoRAs` -- `Scan LoRAs` -- `Select All Compatible` -- `Detected LoRAs (toggle enabled)` -- `LoRAnado blacklist` +### How "ControlNet Unit 0" Works in Forge Neo -### Detection behavior +In Forge Neo, ControlNet units are 0-indexed under the hood: +- **Unit 0** corresponds to the **1st ControlNet tab/accordion slot** in Forge Neo's ControlNet interface. +- When **`Use Image for ControlNet (Unit 0)`** is enabled, RanbooruX automatically fetches the target booru image and populates Unit 0's control image slot before triggering generation. -Detection prefers strict compatibility signals: +### Understanding & Customizing Anima Settings -1. Filename token matches (word-boundary aware): - - PonyXL: `pony`, `pony xl`, `pony-diffusion`, `ponydiffusion`, `pdxl`, `xlp` - - Anima: `anima` -2. Metadata matches from relevant base-model/architecture keys only - - avoids scanning unrelated metadata fields that previously caused false positives +When an Anima model is loaded and **`Auto-detect Anima model`** is enabled: +- **Tag Formatting**: Automatically converts underscores (`_`) to spaces (e.g. `blue_hair` → `blue hair`) for Anima's Qwen3 text encoder. +- **Default Quality Prefix**: Auto-prepends `masterpiece, best quality, score_7, safe, ` if no quality tags are present. +- **Default Negative Prompt**: Auto-fills default negative prompt (`worst quality, low quality, score_1, score_2...`) if negative prompt is empty. +- **Customization**: Uncheck **`Auto-detect Anima model`** to bypass default quality prefixes and negative prompts for 100% custom prompt construction. -If no compatible LoRAs are detected, RanbooruX falls back to all LoRAs in the selected folder so generation is still usable. +### How to Control Anima Sampler Tuning -## Two-pass Img2Img + ADetailer notes +RanbooruX includes dedicated UI controls for Anima sampler and step optimization: -For Img2Img workflows, RanbooruX runs an initial pass, then a dedicated Img2Img pass, then optional manual ADetailer processing. +- **`Auto-tune Img2Img parameters for Anima`** (`anima_tune_img2img`, default ON): + - **When Enabled**: Automatically optimizes step counts, CFG scale (3.0–6.0), and denoising strength (capped at 0.5) tuned for Anima's flow-matching scheduler during Img2Img passes. + - **When Disabled**: RanbooruX preserves your manual step count, CFG scale, and denoising strength set in Forge Neo, giving full manual control to users who prefer custom sampler settings. -> [!NOTE] -> Img2Img is currently **not tested with Anima models/LoRAs**. +### Recommended Settings +- CFG: 4–5 +- Steps: 30–50 +- Sampler: Euler a or er_sde +- Resolution: 512²–1536² +- Clip Skip: 1 -Important behavior: +## RanbooruX vs Original Ranbooru -- first-pass previews are suppressed until final images are ready (preview guard) -- final results are forced back into processed image state for extension/UI consistency -- ADetailer integration uses guarded manual execution to reduce script collisions +Original Ranbooru was a monolithic single-script extension (~1.1k lines). RanbooruX is a complete overhaul built specifically for Forge Neo: + +| Aspect | Original Ranbooru | RanbooruX | +| --- | --- | --- | +| **Target Platform** | Legacy SD WebUI / A1111 | Exclusively **Forge Neo** & **ADetailer Neo** | +| **Architecture** | Single file (`scripts/ranbooru.py`) | Modular package (`ranboorux/`) + script wrappers | +| **Anima Model Support** | None | Full auto-detection, quality defaults, working Img2Img & ControlNet (LLLite) | +| **ADetailer Integration** | None / basic script calling | Guarded two-pass runner supporting ADetailer & ADetailer Neo | +| **Tag Processing** | Ad-hoc string replacements | Bundled Danbooru Tag Catalog (`data/catalogs/danbooru_tags.csv`) | +| **Testing & Quality** | No tests | Complete `pytest` test suite, `mypy`, `ruff`, `black` & CI | +| **Dependency Management** | Implicit / unmanaged | Automated via `requirements.txt` & `install.py` | + +## Forge Neo Technical Notes + +- Target Platform: Developed and tested **strictly for Forge Neo only**. Other WebUI distributions are not supported or tested. +- ControlNet integration is designed for Forge Neo and tested only in that environment. +- Deepbooru support has been removed in RanbooruX. +- The previously bundled `scripts/controlnet.py` has been removed; runtime integration dynamically resolves external/builtin ControlNet paths. +- InputAccordion includes compatibility fallbacks for environments where it is unavailable. +- Gradio update calls are routed through compatibility helpers for Gradio 3/4 behavior. + +## Developer & Verification Guide + +RanbooruX uses a modular architecture with comprehensive automated tests, linting, and type safety checks. + +### Repository Architecture +- `scripts/ranbooru.py`: WebUI Extension entry point and Gradio UI definition. +- `ranboorux/`: Core modular package (catalog pipeline, booru API clients, ADetailer runtime/orchestration, ControlNet integration, Anima model detection). +- `tests/`: Automated test suite covering wrappers, catalog processing, ADetailer runtime, and lifecycle contracts. +- `tools/`: CI helper scripts (`check_no_gradio_update.py`, `repo_guard.py`). + +### Cross-Platform Development Commands + +Run tests, linters, and type checkers locally in your operating system environment: + +#### Windows (PowerShell) +```powershell +$env:PYTHONPATH="." +python -m pytest tests/ -q +python -m ruff check scripts/ranbooru.py ranboorux tests tools install.py +python -m black --check scripts/ranbooru.py ranboorux tests tools install.py +python -m mypy ranboorux --warn-return-any --warn-unused-ignores +``` + +#### Windows (Command Prompt `cmd.exe`) +```cmd +set PYTHONPATH=. +python -m pytest tests/ -q +python -m ruff check scripts/ranbooru.py ranboorux tests tools install.py +python -m black --check scripts/ranbooru.py ranboorux tests tools install.py +python -m mypy ranboorux --warn-return-any --warn-unused-ignores +``` + +#### Linux / macOS (`bash` / `zsh`) +```bash +PYTHONPATH=. python3 -m pytest tests/ -q +PYTHONPATH=. python3 -m ruff check scripts/ranbooru.py ranboorux tests tools install.py +PYTHONPATH=. python3 -m black --check scripts/ranbooru.py ranboorux tests tools install.py +PYTHONPATH=. python3 -m mypy ranboorux --warn-return-any --warn-unused-ignores +``` -## RanbooruX vs Original Ranbooru +## LoRAnado (PonyXL & Anima detection) + +> [!NOTE] +> LoRAnado is a legacy feature inherited from original Ranbooru. + +LoRAnado includes detection and control surfaces to reduce incompatible LoRA picks in PonyXL and Anima workflows. + +Controls: +- `Auto-detect PonyXL/Anima-compatible LoRAs` +- `Scan LoRAs` +- `Select All Compatible` +- `Detected LoRAs (toggle enabled)` +- `LoRAnado blacklist` -- Project scope: original Ranbooru is mostly a single-script extension; RanbooruX adds a modular package (`ranboorux/`), a full `tests/` suite, CI/pre-commit/tooling config, and contributor/testing docs. -- Core implementation: `scripts/ranbooru.py` is heavily expanded/refactored with compatibility wrappers and integration boundaries for Forge Neo. -- Feature set: RanbooruX adds Danbooru tag-catalog processing (bundled/custom CSV + validation/import), `Quick Strip`, richer removal filters, and a diagnostics panel. -- Integration flow: RanbooruX hardens Img2Img + ControlNet + ADetailer behavior on Forge Neo with safer two-pass processing and guarded/manual ADetailer execution. -- LoRAnado: RanbooruX introduces PonyXL & Anima-aware LoRA detection/selection controls and blacklist support. -- Deepbooru Removal: Deepbooru support has been removed in RanbooruX. +Detection matches PonyXL and Anima model signatures based on filename tokens and model metadata keys. If no compatible LoRAs are detected, RanbooruX falls back to all LoRAs in the target directory. ## Credits diff --git a/pics/filters.jpg b/pics/filters.jpg index 52e14f6..035c6ce 100644 Binary files a/pics/filters.jpg and b/pics/filters.jpg differ diff --git a/pics/image.jpg b/pics/image.jpg new file mode 100644 index 0000000..93ebaff Binary files /dev/null and b/pics/image.jpg differ diff --git a/pics/image.png b/pics/image.png deleted file mode 100644 index e3bdc7e..0000000 Binary files a/pics/image.png and /dev/null differ diff --git a/ranboorux/anima_detect.py b/ranboorux/anima_detect.py new file mode 100644 index 0000000..9e5ad2b --- /dev/null +++ b/ranboorux/anima_detect.py @@ -0,0 +1,69 @@ +"""Anima model detection for Forge Neo. + +Provides standalone detection of Anima (2B DiT) models by inspecting +the loaded sd_model object. No dependency on ``modules.shared`` or +``scripts.ranbooru`` — purely parameter-based. +""" + +from __future__ import annotations + +from typing import Any, Optional + + +def _resolve_checkpoint_name(sd_model: Any) -> Optional[str]: + """Return the checkpoint filename from *sd_model* if available.""" + for attr in ("sd_model_checkpoint", "checkpoint", "model_checkpoint"): + value = getattr(sd_model, attr, None) + if value is not None: + return str(value) + return None + + +def get_anima_model_info(sd_model: Any) -> dict[str, Any]: + """Detect whether *sd_model* is an Anima model and return details. + + Returns a dict with keys: + + ``detected`` + ``True`` if the model is identified as Anima. + ``method`` + ``"filename"`` / ``"class_name"`` / ``"none"``. + ``model_name`` + The matched checkpoint filename or class name, or ``""``. + """ + if sd_model is None: + return {"detected": False, "method": "none", "model_name": ""} + + # PRIMARY: checkpoint filename contains "anima" (case-insensitive) + checkpoint = _resolve_checkpoint_name(sd_model) + if checkpoint: + if "anima" in checkpoint.lower(): + return { + "detected": True, + "method": "filename", + "model_name": checkpoint, + } + + # SECONDARY: class name contains "Anima" + class_name = type(sd_model).__name__ + if "Anima" in class_name: + return { + "detected": True, + "method": "class_name", + "model_name": class_name, + } + + return {"detected": False, "method": "none", "model_name": ""} + + +def is_anima_model(sd_model: Any) -> bool: + """Return ``True`` if *sd_model* is an Anima (2B DiT) model. + + Detection order (whichever matches first wins): + + 1. Checkpoint filename containing "anima" (case-insensitive). + 2. Class name containing ``"Anima"`` (e.g. ``class Anima(ForgeDiffusionEngine)``). + + Returns ``False`` for ``None`` input or when no signal is found. + """ + return bool(get_anima_model_info(sd_model)["detected"]) diff --git a/ranboorux/boorus/__init__.py b/ranboorux/boorus/__init__.py new file mode 100644 index 0000000..23566dd --- /dev/null +++ b/ranboorux/boorus/__init__.py @@ -0,0 +1,172 @@ +"""Booru base class and factory function.""" + +import time + +from ranboorux import http_client as rb_http_client + + +class Booru: + def __init__(self, booru_name, base_api_url, http_client=None): + from scripts.ranbooru import Script + + self.booru_name = booru_name + self.base_api_url = base_api_url + self.http = http_client or rb_http_client.BooruSession() + self.headers = {"user-agent": f"Ranbooru Extension/{Script.version} for Forge"} + + def _fetch_data(self, query_url): + from scripts.ranbooru import BooruError, _log + + _log(f"Querying {self.booru_name}: {rb_http_client.redact_url(query_url)}") + max_retries = 3 + for attempt in range(max_retries): + try: + return self.http.get_json(query_url, headers=self.headers, timeout=30) + except Exception as e: + from requests.exceptions import HTTPError, RequestException + + if isinstance(e, HTTPError): + status = getattr(e.response, "status_code", 0) if hasattr(e, "response") else 0 + if status and 400 <= status < 500: + message = rb_http_client.safe_exception_message( + f"fetching data from {self.booru_name}", query_url, e + ) + _log(f"Error {message}") + raise BooruError(f"HTTP Error {message}") from e + elif not isinstance(e, RequestException): + message = rb_http_client.safe_exception_message( + f"fetching data from {self.booru_name}", query_url, e + ) + _log(f"Error {message}") + raise BooruError(f"HTTP Error {message}") from e + + if attempt < max_retries - 1: + sleep_time = 2**attempt + message = rb_http_client.safe_exception_message( + f"fetching data from {self.booru_name}", query_url, e + ) + _log(f"[R] Retry {attempt + 1}/{max_retries} after {sleep_time}s: {message}") + time.sleep(sleep_time) + else: + message = rb_http_client.safe_exception_message( + f"fetching data from {self.booru_name}", query_url, e + ) + _log(f"Error {message}") + raise BooruError(f"HTTP Error {message}") from e + + def _is_direct_image_url(self, url): + """Check if URL is a direct image URL (not from external sites like Pixiv/Twitter)""" + if not url or not isinstance(url, str): + return False + + # Skip external sites that don't provide direct image access + external_sites = [ + "pixiv.net", + "pximg.net", + "twitter.com", + "x.com", + "t.co", + "deviantart.com", + "artstation.com", + "instagram.com", + "facebook.com", + "patreon.com", + "fanbox.cc", + ] + + url_lower = url.lower() + for site in external_sites: + if site in url_lower: + return False + + # Check if URL ends with common image extensions + image_extensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff"] + if any(url_lower.endswith(ext) for ext in image_extensions): + return True + + # Check if URL contains image-serving patterns + if any(pattern in url_lower for pattern in ["/images/", "/img/", "/media/", "/files/"]): + return True + + return False + + def _standardize_post(self, post_data): + from scripts.ranbooru import _split_tag_string, _split_tag_string_override + + post = {} + # extract tags in a robust way; some APIs return categorized tags as dicts + raw_tags = post_data.get("tags", post_data.get("tag_string", "")) + # store categorized lists when possible + artist_tags = [] + character_tags = [] + copyright_tags = [] + if isinstance(post_data.get("tags"), dict): + tags_dict = post_data.get("tags") + # e621 style: tags dict with sublevels + if isinstance(tags_dict.get("artist"), list): + artist_tags = tags_dict.get("artist", []) + if isinstance(tags_dict.get("character"), list): + character_tags = tags_dict.get("character", []) + if isinstance(tags_dict.get("copyright"), list): + copyright_tags = tags_dict.get("copyright", []) + if "tag_string_artist" in post_data: + parsed = _split_tag_string_override(post_data.get("tag_string_artist")) + if parsed is not None: + artist_tags = parsed + if "tag_string_character" in post_data: + parsed = _split_tag_string_override(post_data.get("tag_string_character")) + if parsed is not None: + character_tags = parsed + if "tag_string_copyright" in post_data: + parsed = _split_tag_string_override(post_data.get("tag_string_copyright")) + if parsed is not None: + copyright_tags = parsed + + # For boorus that don't provide categorized tags, try to extract character tags from the main tag string + # This handles cases like Gelbooru/Danbooru where character tags are mixed with other tags + if not character_tags and isinstance(raw_tags, str): + all_tags = _split_tag_string(raw_tags) + for tag in all_tags: + # Common patterns for character tags: contains parentheses (series name) or ends with specific patterns + if ( + ("(" in tag and ")" in tag) + or tag.endswith("_(series)") + or tag.endswith("_(character)") + ): + character_tags.append(tag) + # Also catch some common character name patterns (this is heuristic but should catch most) + elif any( + series in tag.lower() + for series in [ + "genshin_impact", + "touhou", + "fate_", + "azur_lane", + "kantai_collection", + "pokemon", + ] + ): + character_tags.append(tag) + + post["tags"] = raw_tags + post["artist_tags"] = artist_tags + post["character_tags"] = character_tags + post["copyright_tags"] = copyright_tags + post["score"] = post_data.get("score", 0) + post["file_url"] = post_data.get("file_url") + if post["file_url"] is None: + post["file_url"] = post_data.get("large_file_url") + if post["file_url"] is None: + # Check if source is a direct image URL before using it + source_url = post_data.get("source") + if source_url and self._is_direct_image_url(source_url): + post["file_url"] = source_url + else: + post["file_url"] = None + post["id"] = post_data.get("id") + post["rating"] = post_data.get("rating") + post["booru_name"] = self.booru_name + return post + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + raise NotImplementedError diff --git a/ranboorux/boorus/gelbooru.py b/ranboorux/boorus/gelbooru.py new file mode 100644 index 0000000..36aef79 --- /dev/null +++ b/ranboorux/boorus/gelbooru.py @@ -0,0 +1,276 @@ +"""Gelbooru and GelbooruCompatible booru classes.""" + +import random +import time +import xml.etree.ElementTree as ET +from typing import Dict, List, Optional, Tuple +from urllib.parse import quote_plus + +from ranboorux import http_client as rb_http_client +from ranboorux.boorus import Booru + + +class Gelbooru(Booru): + def __init__(self, fringe_benefits, credentials: Optional[Dict[str, str]] = None): + from scripts.ranbooru import POST_AMOUNT, _sanitize_gelbooru_credential + + super().__init__( + "Gelbooru", + f"https://gelbooru.com/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", + ) + self.fringeBenefits = fringe_benefits + credentials = credentials or {} + self.api_key = ( + _sanitize_gelbooru_credential(credentials.get("api_key")) + if isinstance(credentials, dict) + else "" + ) + self.user_id = ( + _sanitize_gelbooru_credential(credentials.get("user_id")) + if isinstance(credentials, dict) + else "" + ) + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + from scripts.ranbooru import BooruError + + _r.COUNT = 0 + all_fetched_posts = [] + if not self.api_key or not self.user_id: + raise BooruError( + "Gelbooru requires an API key and user ID. Set them under RanbooruX \u00bb Gelbooru settings." + ) + credentials_query = ( + f"&api_key={quote_plus(self.api_key)}&user_id={quote_plus(self.user_id)}" + ) + if post_id: + query_url = f"{self.base_api_url}{credentials_query}&id={post_id}{tags_query}" + fetched_data = self._fetch_data(query_url) + if fetched_data and "post" in fetched_data and isinstance(fetched_data["post"], list): + all_fetched_posts = fetched_data["post"] + _r.COUNT = len(all_fetched_posts) + print(f"[R] Found {_r.COUNT} post(s) for ID: {post_id}") + else: + page = random.randint(0, max_pages - 1) + query_url = f"{self.base_api_url}{credentials_query}&pid={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if fetched_data and "post" in fetched_data and isinstance(fetched_data["post"], list): + all_fetched_posts = fetched_data["post"] + if ( + fetched_data + and "@attributes" in fetched_data + and "count" in fetched_data["@attributes"] + ): + try: + _r.COUNT = int(fetched_data["@attributes"]["count"]) + except Exception: + _r.COUNT = len(all_fetched_posts) + else: + _r.COUNT = len(all_fetched_posts) + print( + f"[R] Fetched {len(all_fetched_posts)} posts from page {page}. Reported total (approx): {_r.COUNT}" + ) + return [self._standardize_post(post) for post in all_fetched_posts] + + +class GelbooruCompatible(Booru): + RETRIABLE_STATUS = {429, 500, 502, 503, 504} + + def __init__( + self, base_url: str, retries: int = 3, backoff: float = 1.5, log_diagnostics: bool = True + ): + from scripts.ranbooru import _sanitize_gelbooru_compat_base_url + + sanitized = _sanitize_gelbooru_compat_base_url(base_url) + if not sanitized: + raise ValueError("Invalid Gelbooru-compatible base URL.") + self.base_url = sanitized + self.retries = max(1, retries) + self.backoff = max(0.5, backoff) + self.log_diagnostics = log_diagnostics + self._post_endpoint = f"{self.base_url}/index.php?page=dapi&s=post&q=index" + self._tag_endpoint = f"{self.base_url}/index.php?page=dapi&s=tag&q=index" + self._alias_endpoint = f"{self.base_url}/index.php?page=dapi&s=tag_alias&q=index" + super().__init__("Gelbooru-Compatible", self._post_endpoint) + + def _perform_request(self, url: str): + from scripts.ranbooru import BooruError + + last_error: Optional[Exception] = None + for attempt in range(1, self.retries + 1): + try: + response = self.http.get(url, headers=self.headers, timeout=30, stream=True) + except Exception as exc: + last_error = exc + self._log_retry(url, attempt, f"Request error: {exc.__class__.__name__}") + else: + if response.status_code in self.RETRIABLE_STATUS: + last_error = BooruError(f"Status {response.status_code}") + self._log_retry(url, attempt, f"Status {response.status_code}") + close = getattr(response, "close", None) + if callable(close): + close() + else: + content = self.http._read_bounded_response( + response, + url, + rb_http_client.DEFAULT_API_MAX_BYTES, + ) + return rb_http_client.BoundedResponse( + url=str(getattr(response, "url", url) or url), + status_code=int(getattr(response, "status_code", 200) or 200), + headers=getattr(response, "headers", {}) or {}, + content=content, + encoding=getattr(response, "encoding", None), + ) + time.sleep(min(self.backoff * attempt, 5.0)) + if last_error is None: + error_summary = "unknown error" + elif isinstance(last_error, BooruError): + error_summary = str(last_error) + else: + error_summary = last_error.__class__.__name__ + raise BooruError( + f"HTTP Error fetching from {self.booru_name}: {error_summary} for {rb_http_client.redact_url(url)}" + ) + + def _log_retry(self, url: str, attempt: int, message: str) -> None: + from scripts.ranbooru import _log + + _log(f"{self.booru_name}: retry {attempt} for {rb_http_client.redact_url(url)} - {message}") + + def _log_snippet(self, response) -> None: + from scripts.ranbooru import _log + + if not self.log_diagnostics: + return + snippet = response.text.strip().replace("\n", " ")[:200] + _log( + f"{self.booru_name}: {rb_http_client.redact_url(getattr(response, 'url', ''))} -> {snippet}" + ) + + def _parse_json_entities(self, payload, entity_key: str) -> Tuple[List[dict], Optional[int]]: + entries: List[dict] = [] + approx = None + if isinstance(payload, dict): + possible = payload.get(entity_key) + if isinstance(possible, list): + entries = possible + elif isinstance(possible, dict): + entries = [possible] + attrs = payload.get("@attributes") + if isinstance(attrs, dict) and "count" in attrs: + try: + approx = int(attrs["count"]) + except (TypeError, ValueError): + approx = None + elif isinstance(payload, list): + entries = payload + return entries, approx + + def _parse_xml_entities( + self, text_payload: str, entity_key: str + ) -> Tuple[List[dict], Optional[int]]: + from scripts.ranbooru import BooruError + + probe = (text_payload or "").lower() + if (" Tuple[List[dict], int]: + from scripts.ranbooru import BooruError + + json_url = f"{url_base}&json=1" + try: + response = self._perform_request(json_url) + self._log_snippet(response) + ct = (response.headers.get("content-type") or "").lower() + text_head = (response.text or "").lstrip()[:64].lower() + if ( + "html" in ct + or text_head.startswith(" 0 else 0 + query_base = f"{self._post_endpoint}&limit={POST_AMOUNT}&pid={page}{tags_query}" + posts, approx = self._request_dapi(query_base, "post") + _r.COUNT = approx + print( + f"[R] Gelbooru-compatible: fetched {len(posts)} posts from page {page}. Reported count={approx}" + ) + standardized = [] + for post in posts: + normalized = self._standardize_post(post) + normalized["source_base_url"] = self.base_url + standardized.append(normalized) + return standardized + + def get_tags(self, name_pattern: Optional[str] = None, limit: int = 100) -> List[dict]: + query = f"{self._tag_endpoint}&limit={limit}" + if name_pattern: + query += f"&name_pattern={quote_plus(name_pattern)}" + tags, _ = self._request_dapi(query, "tag") + return tags + + def get_tag_aliases(self, name_pattern: Optional[str] = None, limit: int = 100) -> List[dict]: + query = f"{self._alias_endpoint}&limit={limit}" + if name_pattern: + query += f"&name_pattern={quote_plus(name_pattern)}" + aliases, _ = self._request_dapi(query, "tag_alias") + return aliases diff --git a/ranboorux/boorus/simple.py b/ranboorux/boorus/simple.py new file mode 100644 index 0000000..dba20da --- /dev/null +++ b/ranboorux/boorus/simple.py @@ -0,0 +1,267 @@ +"""Config-driven booru subclasses for 8 simple booru APIs. + +Each subclass has a unique base_url and slight variations in get_posts(). +""" + +import random + +from ranboorux.boorus import Booru + + +class Danbooru(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__("Danbooru", f"https://danbooru.donmai.us/posts.json?limit={POST_AMOUNT}") + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + query_url = f"https://danbooru.donmai.us/posts/{post_id}.json" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, dict) and "id" in fetched_data: + all_fetched_posts = [fetched_data] + _r.COUNT = len(all_fetched_posts) + print(f"[R] Found {_r.COUNT} post(s) for ID: {post_id}") + else: + page = random.randint(1, max_pages) + query_url = f"{self.base_api_url}&page={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, list): + all_fetched_posts = fetched_data + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from page {page}.") + return [self._standardize_post(post) for post in all_fetched_posts if post] + + +class XBooru(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__( + "XBooru", + f"https://xbooru.com/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", + ) + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + query_url = f"{self.base_api_url}&id={post_id}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, dict) and "id" in fetched_data: + all_fetched_posts = [fetched_data] + else: + page = random.randint(0, max_pages - 1) + query_url = f"{self.base_api_url}&pid={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, list): + all_fetched_posts = fetched_data + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from XBooru.") + standardized_posts = [] + for post_data in all_fetched_posts: + post = self._standardize_post(post_data) + if "directory" in post_data and "image" in post_data: + post["file_url"] = ( + f"https://xbooru.com/images/{post_data['directory']}/{post_data['image']}" + ) + standardized_posts.append(post) + return standardized_posts + + +class Rule34(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__( + "Rule34", + f"https://api.rule34.xxx/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", + ) + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + query_url = f"{self.base_api_url}&id={post_id}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, dict) and "id" in fetched_data: + all_fetched_posts = [fetched_data] + else: + page = random.randint(0, max_pages - 1) + query_url = f"{self.base_api_url}&pid={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, list): + all_fetched_posts = fetched_data + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from Rule34.") + return [self._standardize_post(post) for post in all_fetched_posts] + + +class Safebooru(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__( + "Safebooru", + f"https://safebooru.org/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", + ) + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + query_url = f"{self.base_api_url}&id={post_id}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, dict) and "id" in fetched_data: + all_fetched_posts = [fetched_data] + else: + page = random.randint(0, max_pages - 1) + query_url = f"{self.base_api_url}&pid={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, list): + all_fetched_posts = fetched_data + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from Safebooru.") + standardized_posts = [] + for post_data in all_fetched_posts: + post = self._standardize_post(post_data) + if "directory" in post_data and "image" in post_data: + post["file_url"] = ( + f"https://safebooru.org/images/{post_data['directory']}/{post_data['image']}" + ) + standardized_posts.append(post) + return standardized_posts + + +class Konachan(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__("Konachan", f"https://konachan.com/post.json?limit={POST_AMOUNT}") + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + print("[R] Warn: Konachan does not support post IDs.") + return [] + page = random.randint(1, max_pages) + query_url = f"{self.base_api_url}&page={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, list): + all_fetched_posts = fetched_data + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from Konachan.") + return [self._standardize_post(post) for post in all_fetched_posts] + + +class Yandere(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__("Yandere", f"https://yande.re/post.json?limit={POST_AMOUNT}") + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + print("[R] Warn: Yandere does not support post IDs.") + return [] + page = random.randint(1, max_pages) + query_url = f"{self.base_api_url}&page={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, list): + all_fetched_posts = fetched_data + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from Yandere.") + return [self._standardize_post(post) for post in all_fetched_posts] + + +class AIBooru(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__("AIBooru", f"https://aibooru.online/posts.json?limit={POST_AMOUNT}") + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + print("[R] Warn: AIBooru does not support post IDs.") + return [] + page = random.randint(1, max_pages) + query_url = f"{self.base_api_url}&page={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if isinstance(fetched_data, list): + all_fetched_posts = fetched_data + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from AIBooru.") + standardized_posts = [] + for post_data in all_fetched_posts: + post = self._standardize_post(post_data) + post["tags"] = post_data.get("tag_string", "") + standardized_posts.append(post) + return standardized_posts + + +class e621(Booru): + def __init__(self): + from scripts.ranbooru import POST_AMOUNT + + super().__init__("e621", f"https://e621.net/posts.json?limit={POST_AMOUNT}") + + def get_posts(self, tags_query="", max_pages=10, post_id=None): + import scripts.ranbooru as _r + + _r.COUNT = 0 + all_fetched_posts = [] + if post_id: + print("[R] Warn: e621 does not support post IDs.") + return [] + page = random.randint(1, max_pages) + query_url = f"{self.base_api_url}&page={page}{tags_query}" + fetched_data = self._fetch_data(query_url) + if ( + isinstance(fetched_data, dict) + and "posts" in fetched_data + and isinstance(fetched_data["posts"], list) + ): + all_fetched_posts = fetched_data["posts"] + _r.COUNT = len(all_fetched_posts) + print(f"[R] Fetched {_r.COUNT} posts from e621.") + standardized_posts = [] + for post_data in all_fetched_posts: + post = self._standardize_post(post_data) + temp_tags = [] + sublevels = ["general", "artist", "copyright", "character", "species"] + if "tags" in post_data: + for sublevel in sublevels: + if sublevel in post_data["tags"] and isinstance( + post_data["tags"][sublevel], list + ): + temp_tags.extend(post_data["tags"][sublevel]) + post["tags"] = " ".join(temp_tags) + if ( + "score" in post_data + and isinstance(post_data["score"], dict) + and "total" in post_data["score"] + ): + post["score"] = post_data["score"]["total"] + standardized_posts.append(post) + return standardized_posts diff --git a/ranboorux/http_client.py b/ranboorux/http_client.py new file mode 100644 index 0000000..83dc1ff --- /dev/null +++ b/ranboorux/http_client.py @@ -0,0 +1,612 @@ +from __future__ import annotations + +import ipaddress +import json +import re +import socket +from dataclasses import dataclass +from typing import Any, Mapping, Optional +from urllib.parse import parse_qsl, quote, urlencode, urljoin, urlparse + +import requests +from requests.adapters import HTTPAdapter +from urllib3 import PoolManager +from urllib3.connection import HTTPConnection, HTTPSConnection +from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool + +try: + import requests_cache +except Exception: # pragma: no cover - requests_cache is optional in host tests + requests_cache = None + + +SENSITIVE_QUERY_PARAMS = { + "x-amz-credential", + "x-amz-signature", + "x-amz-security-token", + "x-amz-date", + "x-goog-signature", + "x-goog-credential", + "signature", + "sig", + "token", + "access_token", + "authorization", + "key", + "api_key", + "user_id", + "password", +} +REDIRECT_STATUSES = {301, 302, 303, 307, 308} +MAX_REDIRECTS = 5 +STREAM_CHUNK_SIZE = 64 * 1024 +DEFAULT_API_MAX_BYTES = 5 * 1024 * 1024 + + +class ResponseTooLargeError(RuntimeError): + pass + + +class UnsafeUrlError(ValueError): + pass + + +class InvalidContentTypeError(RuntimeError): + pass + + +class BooruResponseError(ValueError): + pass + + +@dataclass +class BoundedResponse: + url: str + status_code: int + headers: Mapping[str, str] + content: bytes + encoding: Optional[str] = None + + @property + def text(self) -> str: + return self.content.decode(self.encoding or "utf-8", errors="replace") + + def json(self) -> Any: + return json.loads(self.content.decode(self.encoding or "utf-8")) + + def raise_for_status(self) -> None: + if 400 <= int(self.status_code) < 600: + raise RuntimeError(f"HTTP status {self.status_code} for {redact_url(self.url)}") + + +def redact_url(url: object) -> str: + text = str(url or "") + if not text: + return text + try: + parsed = urlparse(text) + if not parsed.query: + return text + qsl = parse_qsl(parsed.query, keep_blank_values=True) + new_qsl = [] + for name, value in qsl: + if name.lower() in SENSITIVE_QUERY_PARAMS: + new_qsl.append((name, "")) + else: + new_qsl.append((name, value)) + + query_parts = [] + for name, val in new_qsl: + if val == "": + query_parts.append(f"{quote(name)}=") + else: + query_parts.append(urlencode([(name, val)])) + + return parsed._replace(query="&".join(query_parts)).geturl() + except Exception: + return text + + +def redact_paths(text: str) -> str: + if not text: + return text + + idx = 0 + result = [] + n = len(text) + + def detect_prefix(pos): + # 1. file-URI or file:// + if text[pos:].lower().startswith("file:" + "///"): + return 8, "file" + if text[pos:].lower().startswith("file:" + "//"): + return 7, "file" + + # 2. UNC path starts with \\ + if text[pos:].startswith("\\\\"): + rest = text[pos + 2 :] + if rest and (rest[0].isalnum() or rest[0] in "._-"): + return 2, "unc" + + # 3. Windows drive path: [a-zA-Z]:\ or [a-zA-Z]:/ + is_word_boundary = pos == 0 or not text[pos - 1].isalnum() + if is_word_boundary and pos + 2 < n: + if text[pos].isalpha() and text[pos + 1] == ":" and text[pos + 2] in "\\/": + return 3, "win" + + # 4. POSIX absolute path: starts with / and not followed by / + is_posix_boundary = pos == 0 or (not text[pos - 1].isalnum() and text[pos - 1] != "/") + if is_posix_boundary and text[pos] == "/": + if pos + 1 < n and text[pos + 1] == "/": + return None + return 1, "posix" + + return None + + while idx < n: + prefix_info = detect_prefix(idx) + if prefix_info is None: + result.append(text[idx]) + idx += 1 + continue + + prefix_len, ptype = prefix_info + start_path_idx = idx + + scan_idx = idx + prefix_len + bracket_stack = [] + + while scan_idx < n: + char = text[scan_idx] + + # Stop on quotes, tabs, newlines + if char in "'\"`\t\r\n": + break + + # Stop on unmatched brackets + if char in "([{": + bracket_stack.append(char) + elif char in ")]}": + if not bracket_stack: + break + top = bracket_stack.pop() + if ( + (char == ")" and top != "(") + or (char == "]" and top != "[") + or (char == "}" and top != "{") + ): + break + + # Stop on trailing punctuation followed by space or end of string + is_last = scan_idx + 1 == n + next_char = text[scan_idx + 1] if not is_last else "" + if char in ".,!?;" and (is_last or next_char.isspace()): + break + + # Stop before another path prefix or a URL starts + rem = text[scan_idx:] + if rem.lower().startswith(("http://", "https://", "file:" + "//")): + break + + is_new_path_boundary = scan_idx > 0 and text[scan_idx - 1] not in "\\/:" + if is_new_path_boundary and rem.startswith("\\\\"): + break + + is_rem_word_boundary = scan_idx == 0 or not text[scan_idx - 1].isalnum() + if is_new_path_boundary and is_rem_word_boundary and len(rem) >= 3: + if rem[0].isalpha() and rem[1] == ":" and rem[2] in "\\/": + break + if ( + is_new_path_boundary + and is_rem_word_boundary + and rem.startswith("/") + and not rem.startswith("//") + ): + break + + scan_idx += 1 + + path_str = text[start_path_idx:scan_idx] + stripped_path = path_str.rstrip() + trailing_spaces = path_str[len(stripped_path) :] + + result.append("") + result.append(trailing_spaces) + idx = scan_idx + + return "".join(result) + + +def redact_urls_in_text(text: str) -> str: + # Find all http/https URLs in the text + url_pattern = re.compile(r"https?://[^\s'\")]+", re.IGNORECASE) + + def repl(match): + return redact_url(match.group(0)) + + return url_pattern.sub(repl, text) + + +def sanitize_exception_text(text: str) -> str: + if not text: + return text + # First, redact paths (so URL-like file-URI paths get redacted completely) + text = redact_paths(text) + # Next, redact any remaining HTTP/HTTPS URLs + text = redact_urls_in_text(text) + return text + + +def sanitize_exception(exc: Exception) -> Exception: + if isinstance( + exc, + ( + UnsafeUrlError, + ResponseTooLargeError, + InvalidContentTypeError, + BooruResponseError, + ), + ): + return exc + return RuntimeError(sanitize_exception_text(str(exc))) + + +def safe_exception_message(operation: str, url: object, exc: BaseException) -> str: + sanitized_msg = sanitize_exception_text(str(exc)) + return f"{operation} failed for {redact_url(url)} ({exc.__class__.__name__}: {sanitized_msg})" + + +def _has_sensitive_query(url: object) -> bool: + text = str(url or "") + if not text: + return False + try: + parsed = urlparse(text) + if not parsed.query: + return False + qsl = parse_qsl(parsed.query, keep_blank_values=True) + for name, _ in qsl: + if name.lower() in SENSITIVE_QUERY_PARAMS: + return True + except Exception: + pass + return False + + +def _is_public_ip(address: object) -> bool: + try: + parsed = ipaddress.ip_address(str(address)) + except ValueError: + return False + return bool( + parsed.is_global + and not parsed.is_loopback + and not parsed.is_link_local + and not parsed.is_multicast + and not parsed.is_unspecified + and not parsed.is_reserved + ) + + +def _close_socket(sock: object) -> None: + close = getattr(sock, "close", None) + if callable(close): + close() + + +def _validate_connected_socket(sock: object) -> None: + getpeername = getattr(sock, "getpeername", None) + if not callable(getpeername): + return + peer = getpeername() + address = peer[0] if isinstance(peer, tuple) and peer else None + if address is None: + raise UnsafeUrlError("Connected socket has no peer address") + if not _is_public_ip(address): + _close_socket(sock) + raise UnsafeUrlError(f"Connected peer resolves to a blocked address: {address}") + + +class _SafeHTTPConnection(HTTPConnection): + def _new_conn(self): + sock = super()._new_conn() + _validate_connected_socket(sock) + return sock + + +class _SafeHTTPSConnection(HTTPSConnection): + def _new_conn(self): + sock = super()._new_conn() + _validate_connected_socket(sock) + return sock + + +class _SafeHTTPConnectionPool(HTTPConnectionPool): + ConnectionCls = _SafeHTTPConnection + + +class _SafeHTTPSConnectionPool(HTTPSConnectionPool): + ConnectionCls = _SafeHTTPSConnection + + +class _SafeHTTPAdapter(HTTPAdapter): + def init_poolmanager(self, connections, maxsize, block=False, **pool_kwargs): + self.poolmanager = PoolManager( + num_pools=connections, + maxsize=maxsize, + block=block, + **pool_kwargs, + ) + self.poolmanager.pool_classes_by_scheme = { + "http": _SafeHTTPConnectionPool, + "https": _SafeHTTPSConnectionPool, + } + + def proxy_manager_for(self, proxy, **proxy_kwargs): + manager = super().proxy_manager_for(proxy, **proxy_kwargs) + if hasattr(manager, "pool_classes_by_scheme"): + manager.pool_classes_by_scheme = { + "http": _SafeHTTPConnectionPool, + "https": _SafeHTTPSConnectionPool, + } + return manager + + +def _resolve_host(hostname: str, port: Optional[int]) -> list[str]: + try: + return [str(ipaddress.ip_address(hostname))] + except ValueError: + pass + try: + infos = socket.getaddrinfo(hostname, port, type=socket.SOCK_STREAM) + except OSError as exc: + raise UnsafeUrlError(f"Could not resolve outbound host: {hostname}") from exc + addresses = [] + for info in infos: + sockaddr = info[4] + if sockaddr: + addresses.append(str(sockaddr[0])) + if not addresses: + raise UnsafeUrlError(f"Could not resolve outbound host: {hostname}") + return addresses + + +def validate_outbound_url(url: object) -> str: + text = str(url or "").strip() + parsed = urlparse(text) + if parsed.scheme not in {"http", "https"}: + raise UnsafeUrlError(f"Unsupported outbound URL scheme: {parsed.scheme or ''}") + if not parsed.hostname: + raise UnsafeUrlError("Outbound URL is missing a host") + if parsed.username or parsed.password: + raise UnsafeUrlError("Outbound URL must not contain userinfo credentials") + addresses = _resolve_host(parsed.hostname, parsed.port) + blocked = [address for address in addresses if not _is_public_ip(address)] + if blocked: + raise UnsafeUrlError(f"Outbound URL resolves to a blocked address: {parsed.hostname}") + return text + + +class BooruSession: + def __init__(self, *, use_cache: bool = False, expire_after: int = 3600): + session_factory = getattr(requests, "Session", None) + self._cache_enabled = bool(use_cache) + self._uncached_session = session_factory() if callable(session_factory) else requests + self._install_safe_adapter(self._uncached_session) + if use_cache: + if requests_cache is None: + raise RuntimeError("requests-cache is required when booru request cache is enabled") + cached_session = getattr(requests_cache, "CachedSession", None) + if not callable(cached_session): + raise RuntimeError("requests-cache CachedSession is unavailable") + self._session = cached_session( + "ranbooru_cache", + backend="sqlite", + expire_after=expire_after, + allowable_codes=(200,), + ) + self._install_safe_adapter(self._session) + return + self._session = self._uncached_session + + @staticmethod + def _install_safe_adapter(session: object) -> None: + mount = getattr(session, "mount", None) + if callable(mount): + adapter = _SafeHTTPAdapter() + mount("http://", adapter) + mount("https://", adapter) + + def _session_for_url(self, url: str): + if self._cache_enabled and _has_sensitive_query(url): + return self._uncached_session + return self._session + + def get( + self, + url: str, + *, + headers: Optional[Mapping[str, str]] = None, + timeout: int = 30, + stream: bool = False, + ): + try: + current_url = validate_outbound_url(url) + request_headers = dict(headers or {}) + history_urls = [] + chain_is_sensitive = _has_sensitive_query(current_url) + + for _ in range(MAX_REDIRECTS + 1): + if chain_is_sensitive: + session = self._uncached_session + else: + session = self._session_for_url(current_url) + + history_urls.append((session, current_url)) + + response = session.get( + current_url, + headers=request_headers, + timeout=timeout, + allow_redirects=False, + stream=stream, + ) + status_code = getattr(response, "status_code", None) + if status_code not in REDIRECT_STATUSES: + return response + location = (getattr(response, "headers", {}) or {}).get("location") + + # Remove from cache if the redirect target contains any sensitive queries + if status_code in REDIRECT_STATUSES and location: + redirect_target = urljoin(current_url, location) + if _has_sensitive_query(redirect_target): + chain_is_sensitive = True + + if chain_is_sensitive: + for hist_session, hist_url in history_urls: + delete_fn = getattr(hist_session, "delete", None) + if callable(delete_fn): + try: + delete_fn(hist_url) + except Exception: + pass + + close = getattr(response, "close", None) + if callable(close): + close() + if not location: + return response + current_url = validate_outbound_url(urljoin(current_url, location)) + raise UnsafeUrlError(f"Too many redirects while fetching {redact_url(url)}") + except Exception as exc: + raise sanitize_exception(exc) from None + + def _read_bounded_response(self, response: object, url: str, max_bytes: int) -> bytes: + response_headers = getattr(response, "headers", {}) or {} + content_length = ( + response_headers.get("content-length") if hasattr(response_headers, "get") else None + ) + if content_length: + try: + if int(content_length) > max_bytes: + raise ResponseTooLargeError( + f"Response from {redact_url(url)} exceeded {max_bytes} bytes" + ) + except ValueError: + pass + + chunks: list[bytes] = [] + total = 0 + try: + iter_content = getattr(response, "iter_content", None) + if callable(iter_content): + for chunk in iter_content(chunk_size=STREAM_CHUNK_SIZE): + if not chunk: + continue + total += len(chunk) + if total > max_bytes: + raise ResponseTooLargeError( + f"Response from {redact_url(url)} exceeded {max_bytes} bytes" + ) + chunks.append(chunk) + return b"".join(chunks) + + content = getattr(response, "content", b"") or b"" + if len(content) > max_bytes: + raise ResponseTooLargeError( + f"Response from {redact_url(url)} exceeded {max_bytes} bytes" + ) + return content + finally: + close = getattr(response, "close", None) + if callable(close): + close() + + def get_json( + self, + url: str, + *, + headers: Optional[Mapping[str, str]] = None, + timeout: int = 30, + max_bytes: int = DEFAULT_API_MAX_BYTES, + ) -> Any: + response = self.get(url, headers=headers, timeout=timeout, stream=True) + try: + response.raise_for_status() + response_headers = getattr(response, "headers", {}) or {} + content_type = ( + response_headers.get("content-type", "") if hasattr(response_headers, "get") else "" + ) + normalized_content_type = content_type.lower().split(";", 1)[0].strip() + if normalized_content_type and "json" not in normalized_content_type: + raise InvalidContentTypeError( + f"Response from {redact_url(url)} was not JSON ({content_type})" + ) + content = self._read_bounded_response(response, url, max_bytes) + encoding = getattr(response, "encoding", None) + try: + return json.loads(content.decode(encoding or "utf-8")) + except Exception as exc: + raise BooruResponseError(sanitize_exception_text(str(exc))) from exc + except Exception: + close = getattr(response, "close", None) + if callable(close): + close() + raise + + def get_text( + self, + url: str, + *, + headers: Optional[Mapping[str, str]] = None, + timeout: int = 30, + max_bytes: int = DEFAULT_API_MAX_BYTES, + ) -> BoundedResponse: + response = self.get(url, headers=headers, timeout=timeout, stream=True) + try: + response.raise_for_status() + content = self._read_bounded_response(response, url, max_bytes) + return BoundedResponse( + url=str(getattr(response, "url", url) or url), + status_code=int(getattr(response, "status_code", 200) or 200), + headers=getattr(response, "headers", {}) or {}, + content=content, + encoding=getattr(response, "encoding", None), + ) + except Exception: + close = getattr(response, "close", None) + if callable(close): + close() + raise + + def get_bytes( + self, + url: str, + *, + headers: Optional[Mapping[str, str]] = None, + timeout: int = 30, + max_bytes: int = 25 * 1024 * 1024, + ) -> bytes: + response = self.get(url, headers=headers, timeout=timeout, stream=True) + try: + response.raise_for_status() + response_headers = getattr(response, "headers", {}) or {} + content_type = ( + response_headers.get("content-type", "") if hasattr(response_headers, "get") else "" + ) + if content_type and not content_type.lower().split(";", 1)[0].startswith("image/"): + raise InvalidContentTypeError( + f"Response from {redact_url(url)} was not an image ({content_type})" + ) + return self._read_bounded_response(response, url, max_bytes) + except Exception: + close = getattr(response, "close", None) + if callable(close): + close() + raise + + def close(self) -> None: + for session in (self._session, self._uncached_session): + close = getattr(session, "close", None) + if callable(close): + close() diff --git a/ranboorux/integrations/adetailer_orchestration.py b/ranboorux/integrations/adetailer_orchestration.py new file mode 100644 index 0000000..87746a3 --- /dev/null +++ b/ranboorux/integrations/adetailer_orchestration.py @@ -0,0 +1,783 @@ +"""ADetailer orchestration — lifecycle management extracted from the Script class. + +Phase 3 of the maintainability refactor: encapsulates all ADetailer lifecycle +methods into a single ``AdetailerOrchestrator`` class that the Script instance +delegates to. + +The orchestrator holds a reference to the Script instance (``self._script``) +and accesses Script-owned state (``_adetailer_state``, ``_adetailer_patches``, +``_host_scope``, class-level ``_ranbooru_*`` flags, etc.) through it. +""" + +import logging +import types +from enum import Enum, auto +from typing import Any, List + +from ranboorux import http_client as rb_http_client +from ranboorux.integrations import adetailer_runtime as rb_adetailer_runtime + + +class AdetailerState(Enum): + """Simplified state machine for the ADetailer lifecycle.""" + + IDLE = auto() + """No generation in progress or ADetailer is unblocked.""" + + INITIAL_PASS = auto() + """First pass is running — ADetailer is blocked / guarded.""" + + IMG2IMG_READY = auto() + """Initial pass done; ready for the img2img pass with ADetailer available.""" + + ADETAILER_ACTIVE = auto() + """Manual ADetailer execution is in progress.""" + + DONE = auto() + """Processing complete; guard flags cleared, ready for next generation.""" + + +_logger = logging.getLogger("ranboorux.adetailer_orch") + + +class AdetailerOrchestrator: + """Encapsulates ADetailer lifecycle management for RanbooruX. + + Receives a reference to the owning ``Script`` instance and delegates + Script-owned state access through ``self._script``. + """ + + def __init__(self, script_instance: Any) -> None: + self._script: Any = script_instance + self._state: AdetailerState = AdetailerState.IDLE + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _is_adetailer_script(script: object) -> bool: + """Check if a script is an ADetailer script.""" + try: + if script is None: + return False + script_name = ( + script.__class__.__name__.lower() + if hasattr(script, "__class__") + else str(script).lower() + ) + return ( + "adetailer" in script_name + or "afterdetailer" in script_name + or "after_detailer" in script_name + or "ad_script" in script_name + ) + except Exception as exc: + _logger.warning( + "Failed to inspect ADetailer script type: %s", + rb_http_client.sanitize_exception_text(str(exc)), + ) + return False + + # ------------------------------------------------------------------ + # Queries + # ------------------------------------------------------------------ + + def is_adetailer_enabled(self) -> bool: + """Return whether the Script-level ADetailer support toggle is on.""" + return bool(getattr(self._script, "_adetailer_support_enabled", False)) + + def _is_adetailer_enabled(self) -> bool: + """Alias kept for internal callers during extraction.""" + return self.is_adetailer_enabled() + + # ------------------------------------------------------------------ + # Initial-pass lifecycle + # ------------------------------------------------------------------ + + def _mark_initial_pass(self, p: object) -> None: + """Mark that we are in the initial pass so ADetailer can be intercepted later.""" + self._state = AdetailerState.INITIAL_PASS + try: + print("[R] Marking initial pass - ADetailer will run on img2img results instead") + + # Clear any previous hard-disable flag for ADetailer + try: + if hasattr(p, "_ad_disabled") and getattr(p, "_ad_disabled", False): + self._script._host_scope.set_attr(p, "_ad_disabled", False) + print("[R] Cleared p._ad_disabled from previous generation") + except Exception as _e: + print(f"[R] WARN: Could not clear p._ad_disabled: {_e}") + + # Clear our class-level guard + self._script._set_adetailer_block(False) + self._script._adetailer_state.initial_pass_suppressed = False + # Clear pipeline-level guard flag + setattr(self._script.__class__, "_ranbooru_block_all_adetailer", False) + + # Install runner guard (idempotent) + self._install_scriptrunner_guard(p) + + # CRITICAL: Re-enable any ADetailer scripts from previous generation + self._reenable_adetailer_from_previous_generation() + + # Just set a flag that we are in initial pass + self._script._ranbooru_initial_pass = True + + # Store reference to processing object for later use + self._script._initial_pass_p = p + + except Exception as e: + print(f"[R] Error marking initial pass: {e}") + + def _early_adetailer_protection(self, p: object) -> None: + """Complete ADetailer blocking during initial pass — remove scripts entirely.""" + if not self.is_adetailer_enabled(): + return + try: + print("[R Process] Early ADetailer protection activated") + + # Check if we are in the initial pass + if getattr(self._script, "_ranbooru_initial_pass", False): + print("[R Process] Detected initial pass - COMPLETELY BLOCKING ADetailer") + + # Set comprehensive block flags + self._script._host_scope.set_attr(p, "_ranbooru_skip_initial_adetailer", True) + self._script._host_scope.set_attr(p, "_ranbooru_suppress_all_processing", True) + self._script._host_scope.set_attr(p, "_ranbooru_initial_pass_only", True) + self._script._host_scope.set_attr(p, "_ad_disabled", True) + self._script._adetailer_state.initial_pass_suppressed = True + + # CRITICAL: Completely remove ADetailer scripts from the runner during initial pass + self._remove_adetailer_from_runner(p) + + # Set multiple block flags to ensure no ADetailer execution + self._script._set_adetailer_block(True) + setattr(self._script.__class__, "_ranbooru_block_all_adetailer", True) + setattr(self._script.__class__, "_adetailer_global_guard_active", True) + self._script._adetailer_state.global_guard_active = True + + print( + "[R Process] ADetailer completely blocked for initial pass " + "- will be restored for manual img2img processing" + ) + self._state = AdetailerState.INITIAL_PASS + + except Exception as e: + print(f"[R Process] Error in early ADetailer protection: {e}") + + def _remove_adetailer_from_runner(self, p: object) -> None: + """Temporarily remove ADetailer scripts from the script runner during initial pass.""" + try: + if not hasattr(p, "scripts") or p.scripts is None: + return + + # Store original scripts for restoration + if not hasattr(self._script, "_stored_adetailer_scripts"): + self._script._stored_adetailer_scripts = {"alwayson": [], "regular": []} + + # Remove ADetailer from alwayson_scripts + if hasattr(p.scripts, "alwayson_scripts") and p.scripts.alwayson_scripts: + original_alwayson = list(p.scripts.alwayson_scripts) + filtered_alwayson = [ + s for s in original_alwayson if not self._is_adetailer_script(s) + ] + removed_alwayson = [s for s in original_alwayson if self._is_adetailer_script(s)] + + p.scripts.alwayson_scripts = filtered_alwayson + self._script._stored_adetailer_scripts["alwayson"] = removed_alwayson + print( + f"[R Process] Removed {len(removed_alwayson)} ADetailer scripts " + "from alwayson_scripts" + ) + + # Remove ADetailer from regular scripts + if hasattr(p.scripts, "scripts") and p.scripts.scripts: + original_scripts = list(p.scripts.scripts) + filtered_scripts = [s for s in original_scripts if not self._is_adetailer_script(s)] + removed_scripts = [s for s in original_scripts if self._is_adetailer_script(s)] + + p.scripts.scripts = filtered_scripts + self._script._stored_adetailer_scripts["regular"] = removed_scripts + print(f"[R Process] Removed {len(removed_scripts)} ADetailer scripts from scripts") + + # Also check global script lists (ADetailer-Neo on Forge Neo) + try: + import modules.scripts as scripts_module + + for attr in ("scripts_txt2img", "scripts_img2img"): + global_runner = getattr(scripts_module, attr, None) + if global_runner is None or global_runner is p.scripts: + continue + for list_attr in ("alwayson_scripts", "scripts"): + script_list = getattr(global_runner, list_attr, None) + if not script_list: + continue + adetailer_global = [s for s in script_list if self._is_adetailer_script(s)] + if adetailer_global: + self._script._stored_adetailer_scripts[list_attr] = adetailer_global + print( + f"[R Process] Found {len(adetailer_global)} ADetailer " + f"scripts in global {attr}.{list_attr} " + "(blocking via flags)" + ) + except Exception: + pass + + # Also check global script lists (ADetailer-Neo on Forge Neo) + except Exception as e: + print(f"[R Process] Error removing ADetailer from runner: {e}") + + def _restore_early_adetailer_protection(self, processing_obj: object = None) -> None: + """Restore ADetailer scripts and flags after an interrupted or completed run.""" + try: + self._state = AdetailerState.IMG2IMG_READY + print("[R Process] Restoring ADetailer scripts for manual processing") + + # Clear initial pass/block flags so subsequent generations can run ADetailer + setattr(self._script.__class__, "_ranbooru_block_all_adetailer", False) + setattr(self._script.__class__, "_adetailer_global_guard_active", False) + self._script._set_adetailer_block(False) + + # Determine which processing object's script runner to restore into + candidate_p = ( + processing_obj + or getattr(self._script, "_initial_pass_p", None) + or getattr(self._script, "_current_processing_object", None) + ) + runner = getattr(candidate_p, "scripts", None) if candidate_p else None + + # Restore scripts we removed during the initial pass safeguard + stored = getattr(self._script, "_stored_adetailer_scripts", None) + if stored and runner: + try: + if hasattr(runner, "alwayson_scripts") and stored.get("alwayson"): + for script in stored["alwayson"]: + if script not in runner.alwayson_scripts: + runner.alwayson_scripts.append(script) + print( + f"[R Process] Reattached {len(stored['alwayson'])} " + "ADetailer always-on script(s)" + ) + if hasattr(runner, "scripts") and stored.get("regular"): + for script in stored["regular"]: + if script not in runner.scripts: + runner.scripts.append(script) + print( + f"[R Process] Reattached {len(stored['regular'])} " + "ADetailer on-demand script(s)" + ) + finally: + # Clear stored references so we don't duplicate reinsertion + delattr(self._script, "_stored_adetailer_scripts") + + # Ensure any scripts we hard-disabled are re-enabled for the next generation + if hasattr(self._script, "disabled_adetailer_scripts"): + self._reenable_adetailer_from_previous_generation() + + # Clear temporary protection flag if present + if hasattr(self._script, "_temp_disabled_adetailer"): + delattr(self._script, "_temp_disabled_adetailer") + + print("[R Process] Early protection restoration complete") + + except Exception as e: + print(f"[R Process] Error restoring early ADetailer protection: {e}") + + def _prepare_adetailer_for_img2img(self, p: object) -> None: + """Prepare ADetailer to run on img2img results.""" + if not self.is_adetailer_enabled(): + return + try: + print("[R] Preparing ADetailer to run on img2img results") + + # Clear the initial pass flag so ADetailer knows to run normally + self._script._ranbooru_initial_pass = False + + except Exception as e: + print(f"[R] Error preparing ADetailer: {e}") + + # ------------------------------------------------------------------ + # Full restore / native ADetailer re-enablement + # ------------------------------------------------------------------ + + def _restore_native_adetailer_scripts(self, p: object) -> None: + """Ensure native ADetailer scripts resume running when manual support is disabled.""" + try: + if not self._script._adetailer_patches.is_empty(): + self._script._unpatch_manual_adetailer_overrides() + except Exception as exc: + print(f"[R Before] Warn: Could not unpatch manual ADetailer overrides: {exc}") + try: + self._script._set_adetailer_block(False) + except Exception: + pass + setattr(self._script.__class__, "_ranbooru_block_all_adetailer", False) + setattr(self._script.__class__, "_adetailer_global_guard_active", False) + try: + self._restore_early_adetailer_protection(p) + except Exception as exc: + print(f"[R Before] Warn: Could not restore ADetailer runner state: {exc}") + try: + self._reenable_adetailer_from_previous_generation() + except Exception as exc: + print(f"[R Before] Warn: Could not re-enable ADetailer scripts: {exc}") + try: + restored = self._force_enable_adetailer_scripts(p) + except Exception as exc: + print(f"[R Before] Warn: Could not force-enable ADetailer scripts: {exc}") + restored = 0 + if restored: + print( + f"[R Before] Restored {restored} native ADetailer script(s) " + "after manual toggle was disabled" + ) + if hasattr(self._script, "disabled_adetailer_scripts"): + try: + delattr(self._script, "disabled_adetailer_scripts") + except Exception: + pass + guard_present = False + try: + import modules.scripts as scripts_module + + for runner_attr in ("scripts_txt2img", "scripts_img2img"): + runner = getattr(scripts_module, runner_attr, None) + if runner and getattr(runner, "_ranbooru_guard_installed", False): + guard_present = True + break + except Exception: + guard_present = False + if guard_present: + try: + self._script._reset_script_runner_guards() + except Exception as exc: + print(f"[R Before] Warn: Could not reset script runner guards: {exc}") + self._ensure_native_adetailer_enable_flags(p) + if not self._script._native_adetailer_detected(): + try: + import modules.scripts as scripts_module + + if hasattr(scripts_module, "reload_scripts"): + print("[R Before] Reloading scripts to restore native ADetailer") + scripts_module.reload_scripts() + except Exception as exc: + print(f"[R Before] Warn: Could not reload scripts for ADetailer: {exc}") + + def _force_enable_adetailer_scripts(self, processing_obj: object = None) -> int: + """Return the count of ADetailer scripts restored to their original behaviour.""" + try: + import modules.scripts as scripts_module + except Exception as exc: + print(f"[R Before] Warn: Could not access scripts module to restore ADetailer: {exc}") + return 0 + runners: List[object] = [] + for runner_attr in ("scripts_txt2img", "scripts_img2img"): + runner = getattr(scripts_module, runner_attr, None) + if runner: + runners.append(runner) + if ( + processing_obj is not None + and hasattr(processing_obj, "scripts") + and processing_obj.scripts not in runners + ): + runners.append(processing_obj.scripts) + seen_ids: set = set() + restored_count = 0 + for runner in runners: + if runner is None: + continue + for list_attr in ("alwayson_scripts", "scripts"): + script_list = getattr(runner, list_attr, None) + if not script_list: + continue + for script in script_list: + if not script: + continue + script_id = id(script) + if script_id in seen_ids: + continue + seen_ids.add(script_id) + if not self._is_adetailer_script(script): + continue + restored = False + if hasattr(script, "enabled") and script.enabled is False: + script.enabled = True + restored = True + for method_name in ( + "postprocess", + "process", + "process_batch", + "before_process", + "after_process", + ): + backup_name = f"_ranbooru_original_{method_name}" + if hasattr(script, backup_name): + try: + setattr(script, method_name, getattr(script, backup_name)) + except Exception: + pass + try: + delattr(script, backup_name) + except Exception: + pass + restored = True + for attr in ("_ranbooru_disabled_after_manual", "_ranbooru_disabled_source"): + if hasattr(script, attr): + try: + delattr(script, attr) + except Exception: + pass + restored = True + if restored: + restored_count += 1 + if restored_count == 0: + try: + debug_entries = [] + for runner in runners: + if not runner: + continue + for list_attr in ("alwayson_scripts", "scripts"): + script_list = getattr(runner, list_attr, None) + if not script_list: + continue + for script in script_list: + if self._is_adetailer_script(script): + debug_entries.append( + f"{script.__class__.__name__}" + f"(enabled={getattr(script, 'enabled', 'n/a')})" + ) + if debug_entries: + print( + "[R Before] Native ADetailer scripts detected: " + ", ".join(debug_entries) + ) + except Exception: + pass + return restored_count + + def _ensure_native_adetailer_enable_flags(self, processing_obj: Any) -> None: + """Ensure ADetailer enable/skip flags in script_args are set correctly.""" + if not getattr(self._script, "_adetailer_support_enabled", False): + return + try: + args = getattr(processing_obj, "script_args", None) + except Exception as exc: + print(f"[R Before] Native ADetailer: unable to read script_args: {exc}") + return + if not isinstance(args, (list, tuple)) or not args: + print( + "[R Before] Native ADetailer: script_args empty or not list/tuple; " + "skipping flag repair" + ) + return + args_list = list(args) + runners: List[object] = [] + runner = getattr(processing_obj, "scripts", None) + if runner is not None: + runners.append(runner) + try: + import modules.scripts as scripts_module + + for attr in ("scripts_txt2img", "scripts_img2img"): + global_runner = getattr(scripts_module, attr, None) + if global_runner is not None and global_runner not in runners: + runners.append(global_runner) + except Exception as exc: + print(f"[R Before] Native ADetailer: could not gather global runners: {exc}") + candidates: list = [] + for r in runners: + for list_attr in ("alwayson_scripts", "scripts"): + script_list = getattr(r, list_attr, None) + if script_list: + candidates.extend(script_list) + if not candidates: + print("[R Before] Native ADetailer: no script candidates found for flag repair") + return + changed = False + for script in candidates: + if not self._is_adetailer_script(script): + continue + extracted = self._script._extract_adetailer_script_args(script, processing_obj) + sanitized = list(extracted.get("args") or []) + meta = extracted.get("meta") or {} + start_idx = meta.get("slice_start") + end_idx = meta.get("slice_end") + if start_idx is None or end_idx is None: + continue + start_idx = max(0, min(len(args_list), start_idx)) + end_idx = max(start_idx, min(len(args_list), end_idx)) + if not sanitized or end_idx - start_idx != len(sanitized): + slice_view = args_list[start_idx:end_idx] + else: + slice_view = sanitized + print( + f"[R Before] Native ADetailer candidate {script.__class__.__name__} " + f"enabled={getattr(script, 'enabled', 'n/a')} " + f"slice [{start_idx}:{end_idx}] -> {slice_view}" + ) + if not sanitized: + continue + bool_index = 0 + local_changed = False + for offset, val in enumerate(sanitized): + if isinstance(val, bool): + if bool_index == 0 and val is False: + sanitized[offset] = True + local_changed = True + print( + f"[R Before] Set native ADetailer enable flag True at offset {offset}" + ) + elif bool_index == 1 and val is True: + sanitized[offset] = False + local_changed = True + print(f"[R Before] Cleared native ADetailer skip flag at offset {offset}") + bool_index += 1 + elif isinstance(val, dict): + if val.get("ad_tab_enable") is False and val.get("ad_model") not in ( + None, + "", + "None", + ): + val["ad_tab_enable"] = True + local_changed = True + print(f"[R Before] Enabled ad_tab_enable in dict at offset {offset}") + if local_changed: + if end_idx - start_idx == len(sanitized): + args_list[start_idx:end_idx] = sanitized + changed = True + continue + # fallback if lengths mismatch + for offset, val in enumerate(sanitized): + target_idx = start_idx + offset + if target_idx < len(args_list): + args_list[target_idx] = val + else: + args_list.append(val) + changed = True + if changed: + if isinstance(args, list): + processing_obj.script_args = args_list + else: + processing_obj.script_args = tuple(args_list) + print(f"[R Before] Native ADetailer flags updated: {args_list}") + else: + print("[R Before] Native ADetailer flags already enabled; no changes made") + + # ------------------------------------------------------------------ + # Re-enable from previous generation + # ------------------------------------------------------------------ + + def _reenable_adetailer_from_previous_generation(self) -> None: + """Re-enable ALL ADetailer scripts that were disabled in the previous generation.""" + try: + if ( + hasattr(self._script, "disabled_adetailer_scripts") + and self._script.disabled_adetailer_scripts + ): + print( + f"[R] COMPREHENSIVE RE-ENABLE: Restoring " + f"{len(self._script.disabled_adetailer_scripts)} ADetailer script(s) " + "from previous generation" + ) + + for script, original_enabled in self._script.disabled_adetailer_scripts: + source = getattr(script, "_ranbooru_disabled_source", "unknown") + print(f"[R] Re-enabling {script.__class__.__name__} from {source}") + + # Restore original enabled state + if hasattr(script, "enabled"): + script.enabled = original_enabled + + # Restore ALL original methods that were disabled + methods_to_restore = [ + "postprocess", + "process", + "process_batch", + "before_process", + "after_process", + ] + for method_name in methods_to_restore: + original_method_attr = f"_ranbooru_original_{method_name}" + if hasattr(script, original_method_attr): + original_method = getattr(script, original_method_attr) + setattr(script, method_name, original_method) + delattr(script, original_method_attr) + + # Remove our disable flags + if hasattr(script, "_ranbooru_disabled_after_manual"): + delattr(script, "_ranbooru_disabled_after_manual") + if hasattr(script, "_ranbooru_disabled_source"): + delattr(script, "_ranbooru_disabled_source") + + print( + f"[R] COMPREHENSIVE RE-ENABLE: Restored " + f"{len(self._script.disabled_adetailer_scripts)} ADetailer script(s) " + "for new generation" + ) + # Clear the list now that we've re-enabled everything + delattr(self._script, "disabled_adetailer_scripts") + + except Exception as e: + print(f"[R] Error in comprehensive ADetailer re-enable: {e}") + + # ------------------------------------------------------------------ + # Manual ADetailer execution + # ------------------------------------------------------------------ + + def _execute_manual_adetailer(self, p: Any, processed: Any, img2img_results: List[Any]) -> bool: + """Run manual ADetailer on img2img results via the deterministic runtime executor.""" + if not self.is_adetailer_enabled() or not img2img_results: + return False + + self._script._clear_manual_adetailer_skip_flags(p) + adetailer_scripts = rb_adetailer_runtime.gather_adetailer_scripts(p) + if not adetailer_scripts: + print("[R Post] WARN: No ADetailer scripts discovered for manual execution") + return False + + setattr(self._script.__class__, "_ranbooru_manual_adetailer_active", True) + + def build_processed(single_image: object) -> object: + temp_processed = types.SimpleNamespace() + temp_processed.images = [single_image] + temp_processed.image = single_image + for attr in ( + "prompt", + "negative_prompt", + "seed", + "subseed", + "width", + "height", + "cfg_scale", + "steps", + ): + if hasattr(processed, attr): + setattr(temp_processed, attr, getattr(processed, attr)) + return temp_processed + + self._state = AdetailerState.ADETAILER_ACTIVE + try: + result = rb_adetailer_runtime.execute_manual_adetailer( + adetailer_scripts=adetailer_scripts, + images=list(img2img_results), + processing_obj=p, + run_state=self._script._adetailer_state, + patch_registry=self._script._adetailer_patches, + extract_script_args=self._script._extract_adetailer_script_args, + build_processed=build_processed, + isolation_factory=lambda script_obj: self._script._manual_adetailer_script_isolation( + p, + script_obj, + keep_controlnet=self._script._manual_adetailer_requires_controlnet( + self._script._extract_adetailer_script_args(script_obj, p).get("args") or [] + ), + ), + ) + finally: + setattr(self._script.__class__, "_ranbooru_manual_adetailer_active", False) + self._state = AdetailerState.IMG2IMG_READY + + for error in result.errors: + print(f"[R Post] WARN: Manual ADetailer error: {error}") + processed.images.clear() + processed.images.extend(result.images) + img2img_results.clear() + img2img_results.extend(result.images) + if hasattr(p, "processed") and hasattr(p.processed, "images"): + p.processed.images.clear() + p.processed.images.extend(result.images) + return result.successful_processes > 0 + + # ------------------------------------------------------------------ + # ScriptRunner guard + # ------------------------------------------------------------------ + + def _install_scriptrunner_guard(self, p: object) -> None: + """Wrap p.scripts postprocess/postprocess_image to skip ADetailer when blocked.""" + try: + if not hasattr(p, "scripts") or p.scripts is None: + return + runner = p.scripts + if getattr(runner, "_ranbooru_guard_installed", False): + return + rb_adetailer_runtime.install_runner_guard( + runner=runner, + block_flag_fn=lambda: bool( + getattr(self._script.__class__, "_ranbooru_block_all_adetailer", False) + and not getattr( + self._script.__class__, "_ranbooru_manual_adetailer_active", False + ) + ), + patch_registry=self._script._adetailer_patches, + ) + runner._ranbooru_guard_installed = True + self._script._log_patch_event( + "info", "Installed ScriptRunner guard to skip ADetailer when blocked" + ) + except Exception as e: + self._script._log_patch_event("warning", f"Failed to install ScriptRunner guard: {e}") + print(f"[R] Error installing ScriptRunner guard: {e}") + + # ------------------------------------------------------------------ + # Preview guard (shared.state) + # ------------------------------------------------------------------ + + def _install_preview_guard(self) -> None: + """Install a guard around shared.state.assign_current_image to block wrong previews.""" + try: + import modules.shared as shared + + if not hasattr(shared, "state"): + return + state = shared.state + installed_wrapper = getattr(state, "_ranbooru_preview_guard_wrapper", None) + if ( + getattr(state, "_ranbooru_preview_guard_installed", False) + and installed_wrapper is not None + and getattr(state, "assign_current_image", None) is installed_wrapper + ): + self._state = AdetailerState.IMG2IMG_READY + return + if not hasattr(state, "assign_current_image"): + self._state = AdetailerState.IMG2IMG_READY + return + self._state = AdetailerState.INITIAL_PASS + original_assign_current_image = state.assign_current_image + script_class = self._script.__class__ + + def guarded_assign_current_image(img: object) -> Any: + try: + if getattr(script_class, "_ranbooru_preview_guard_on", False): + if getattr(script_class, "_ranbooru_preview_block_all", False): + if not getattr( + script_class, "_ranbooru_preview_block_notice_emitted", False + ): + print( + "[R UI] Preview blocked: withholding intermediary frame " + "until final image is ready" + ) + script_class._ranbooru_preview_block_notice_emitted = True + return + # If we know final dims, only allow those; otherwise block 640x512 + final_dims = getattr(script_class, "_ranbooru_final_dims", None) + if img is not None and hasattr(img, "size"): + if final_dims and img.size != final_dims: + print("[R UI] Preview blocked: mismatched size") + return + if img.size == (640, 512): + print("[R UI] Preview blocked: 640x512 preview") + return + except Exception: + pass + return original_assign_current_image(img) + + self._script._host_scope.patch_attr( + state, "assign_current_image", guarded_assign_current_image + ) + self._script._host_scope.set_attr(state, "_ranbooru_preview_guard_installed", True) + self._script._host_scope.set_attr( + state, "_ranbooru_preview_guard_wrapper", guarded_assign_current_image + ) + print("[R UI] Installed preview guard") + except Exception as e: + print(f"[R UI] Error installing preview guard: {e}") diff --git a/ranboorux/integrations/adetailer_runtime.py b/ranboorux/integrations/adetailer_runtime.py index 2209214..b6cf640 100644 --- a/ranboorux/integrations/adetailer_runtime.py +++ b/ranboorux/integrations/adetailer_runtime.py @@ -359,23 +359,44 @@ def runner_isolation( ) -def _images_differ(original: object, updated: object) -> bool: +def _images_differ(original: object, updated: object, _debug: bool = False) -> bool: if updated is None: + if _debug: + print("[R] _images_differ: updated is None → False") return False if original is None: + if _debug: + print("[R] _images_differ: original is None → True") return True + if original is updated: + if _debug: + print("[R] _images_differ: same object → False") + return False original_size = getattr(original, "size", None) updated_size = getattr(updated, "size", None) if original_size is not None and updated_size is not None and original_size != updated_size: + if _debug: + print(f"[R] _images_differ: size {original_size} vs {updated_size} → True") return True try: original_bytes = original.tobytes() if hasattr(original, "tobytes") else None updated_bytes = updated.tobytes() if hasattr(updated, "tobytes") else None if original_bytes is not None and updated_bytes is not None: - return bool(original_bytes != updated_bytes) + differ = bool(original_bytes != updated_bytes) + if _debug: + o_token = getattr(original, "token", "?") + u_token = getattr(updated, "token", "?") + print( + f"[R] _images_differ: tobytes {o_token} vs {u_token} len={len(original_bytes)}/{len(updated_bytes)} → {differ}" + ) + return differ except Exception: - return original is not updated - return original is not updated + if _debug: + print("[R] _images_differ: tobytes exception → identity check") + return True + if _debug: + print("[R] _images_differ: no tobytes → identity check → True") + return True def _candidate_scripts(adetailer_scripts: Iterable[object]) -> List[object]: @@ -394,12 +415,13 @@ def _candidate_scripts(adetailer_scripts: Iterable[object]) -> List[object]: def _extract_processed_image(temp_processed: object, fallback: object) -> object: - images = getattr(temp_processed, "images", None) - if isinstance(images, list) and images: - return images[0] + # ADetailer modifies pp.image in place; check it BEFORE pp.images image = getattr(temp_processed, "image", None) if image is not None: return image + images = getattr(temp_processed, "images", None) + if isinstance(images, list) and images: + return images[0] return fallback @@ -478,10 +500,22 @@ def execute_manual_adetailer( def gather_adetailer_scripts(processing_obj: object) -> List[object]: - runner = getattr(processing_obj, "scripts", None) - if runner is None: - return [] scripts_list: List[object] = [] - scripts_list.extend(list(getattr(runner, "alwayson_scripts", []) or [])) - scripts_list.extend(list(getattr(runner, "scripts", []) or [])) + + runner = getattr(processing_obj, "scripts", None) + if runner is not None: + scripts_list.extend(list(getattr(runner, "alwayson_scripts", []) or [])) + scripts_list.extend(list(getattr(runner, "scripts", []) or [])) + + try: + import modules.scripts as scripts_module + + for attr in ("scripts_txt2img", "scripts_img2img"): + global_runner = getattr(scripts_module, attr, None) + if global_runner is not None and global_runner is not runner: + scripts_list.extend(list(getattr(global_runner, "alwayson_scripts", []) or [])) + scripts_list.extend(list(getattr(global_runner, "scripts", []) or [])) + except Exception: + pass + return _candidate_scripts(scripts_list) diff --git a/ranboorux/integrations/controlnet.py b/ranboorux/integrations/controlnet.py index 61f67e2..5e5c800 100644 --- a/ranboorux/integrations/controlnet.py +++ b/ranboorux/integrations/controlnet.py @@ -1,12 +1,11 @@ from __future__ import annotations -import importlib import importlib.util import logging import os from types import ModuleType -from ranboorux.requesting import sanitize_exception_text +from ranboorux.http_client import sanitize_exception_text logger = logging.getLogger("ranboorux") diff --git a/ranboorux/loranado.py b/ranboorux/loranado.py index 9c7e8cc..0e4e098 100644 --- a/ranboorux/loranado.py +++ b/ranboorux/loranado.py @@ -11,7 +11,6 @@ def normalize_lora_name(value: object) -> str: text = str(value).strip() if not text: return "" - # Strip extension and return lowercase return os.path.splitext(text)[0].strip().lower() @@ -38,12 +37,10 @@ def filter_candidates( normalize_lora_name(name) for name in blacklist_loras if normalize_lora_name(name) } - # 1. Filter enabled candidates if any are selected filtered = list(candidates) if enabled_selection: filtered = [c for c in filtered if normalize_lora_name(c) in enabled_selection] - # 2. Filter blacklist candidates if blacklist_selection: filtered = [c for c in filtered if normalize_lora_name(c) not in blacklist_selection] @@ -77,7 +74,6 @@ def select_loras( chosen_file = chosen_files[i] lora_name = os.path.splitext(chosen_file)[0] - # Use custom weight if available, otherwise draw randomly if i < len(weights): weight = weights[i] else: diff --git a/ranboorux/mutation_scope.py b/ranboorux/mutation_scope.py new file mode 100644 index 0000000..08b4e72 --- /dev/null +++ b/ranboorux/mutation_scope.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import shutil +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, List, Tuple + +_MISSING = object() +_ANY = object() + + +@dataclass +class RunContext: + temp_paths: List[Path] = field(default_factory=list) + cleanup_errors: List[str] = field(default_factory=list) + + def own_temp_path(self, path: str) -> None: + self.temp_paths.append(Path(path)) + + def cleanup(self) -> None: + for path in reversed(self.temp_paths): + try: + if path.is_dir(): + shutil.rmtree(path, ignore_errors=False) + elif path.exists(): + path.unlink() + except Exception as exc: + self.cleanup_errors.append(f"{path}: {exc}") + self.temp_paths.clear() + + +@dataclass +class HostMutationScope: + context: RunContext = field(default_factory=RunContext) + _snapshots: List[Tuple[object, str, Any, Any]] = field(default_factory=list) + _restored: bool = False + + def snapshot_attr(self, target: object, attr_name: str) -> None: + for existing_target, existing_attr, _value, _expected in self._snapshots: + if existing_target is target and existing_attr == attr_name: + return + value = getattr(target, attr_name, _MISSING) + self._snapshots.append((target, attr_name, value, _ANY)) + + def set_attr(self, target: object, attr_name: str, value: object) -> None: + self.snapshot_attr(target, attr_name) + setattr(target, attr_name, value) + + def patch_attr(self, target: object, attr_name: str, replacement: object) -> None: + for index, (existing_target, existing_attr, value, expected) in enumerate(self._snapshots): + if existing_target is target and existing_attr == attr_name: + if expected is _ANY or getattr(target, attr_name, _MISSING) is expected: + self._snapshots[index] = (target, attr_name, value, replacement) + setattr(target, attr_name, replacement) + return + value = getattr(target, attr_name, _MISSING) + self._snapshots.append((target, attr_name, value, replacement)) + setattr(target, attr_name, replacement) + + def restore(self) -> None: + if self._restored: + return + self._restored = True + for target, attr_name, value, expected in reversed(self._snapshots): + try: + if expected is not _ANY and getattr(target, attr_name, _MISSING) is not expected: + continue + if value is _MISSING: + if hasattr(target, attr_name): + delattr(target, attr_name) + else: + setattr(target, attr_name, value) + except Exception as exc: + self.context.cleanup_errors.append(f"{target!r}.{attr_name}: {exc}") + self._snapshots.clear() + self.context.cleanup() diff --git a/ranboorux/prompting.py b/ranboorux/prompting.py deleted file mode 100644 index 7bd6329..0000000 --- a/ranboorux/prompting.py +++ /dev/null @@ -1,44 +0,0 @@ -from __future__ import annotations - -from typing import Iterable, List, Union - - -def split_prompt_tags(prompt: str) -> List[str]: - if not isinstance(prompt, str): - return [] - return [tag.strip() for tag in prompt.split(",") if tag.strip()] - - -def dedupe_keep_order(tags: Iterable[str]) -> List[str]: - return list(dict.fromkeys(tags)) - - -def remove_repeated_tags(prompt: str) -> str: - tags = split_prompt_tags(prompt) - if not tags: - return "" - return ",".join(dedupe_keep_order(tags)) - - -def limit_prompt_tags(prompt: str, limit_val: Union[int, float, str], mode: str) -> str: - tags = split_prompt_tags(prompt) - if not tags: - return "" - if mode == "Limit": - try: - pct = float(limit_val) - except Exception: - return prompt - if pct <= 0: - return "" - max_count = max(1, int(len(tags) * pct)) - return ",".join(tags[:max_count]) - if mode == "Max": - try: - max_count = int(limit_val) - except Exception: - return prompt - if max_count <= 0: - return prompt - return ",".join(tags[:max_count]) - return prompt diff --git a/ranboorux/run_options.py b/ranboorux/run_options.py index c0bf80c..3c58ccf 100644 --- a/ranboorux/run_options.py +++ b/ranboorux/run_options.py @@ -1,7 +1,8 @@ from __future__ import annotations +# SIZE_OK — cohesive dataclass/schema module; splitting by count scatters related definitions from dataclasses import dataclass -from typing import Dict, Iterable, List, Mapping, Sequence, Tuple +from typing import Any, Dict, List, Mapping, Sequence, Tuple UI_ARGUMENT_FIELDS: Tuple[str, ...] = ( "enabled", @@ -66,6 +67,8 @@ "lora_auto_detect_pony", "lora_detected_loras", "lora_blacklist", + "anima_auto_detect", + "anima_tune_img2img", ) @@ -188,6 +191,8 @@ class RunOptions: lora_auto_detect_pony: object lora_detected_loras: object lora_blacklist: object + anima_auto_detect: bool = True + anima_tune_img2img: bool = True @classmethod def from_script_args(cls, args: Sequence[object]) -> "RunOptions": @@ -195,7 +200,8 @@ def from_script_args(cls, args: Sequence[object]) -> "RunOptions": expected = len(UI_ARGUMENT_FIELDS) if len(values) != expected: raise ValueError(f"Expected {expected} RanbooruX script args, got {len(values)}") - return cls(**dict(zip(UI_ARGUMENT_FIELDS, values))) + kw: Dict[str, Any] = dict(zip(UI_ARGUMENT_FIELDS, values)) + return cls(**kw) def as_dict(self) -> Dict[str, object]: return {field: getattr(self, field) for field in UI_ARGUMENT_FIELDS} @@ -276,9 +282,3 @@ def script_args(self) -> List[object]: if missing: raise ValueError(f"Missing RanbooruX components: {', '.join(missing)}") return [self.components[field] for field in UI_ARGUMENT_FIELDS] - - -def assert_known_fields(fields: Iterable[str]) -> None: - incoming = tuple(fields) - if incoming != UI_ARGUMENT_FIELDS: - raise ValueError("RanbooruX UI argument fields do not match the authoritative order") diff --git a/ranboorux/tag_pipeline.py b/ranboorux/tag_pipeline.py index 4aba5d7..6ae3c0d 100644 --- a/ranboorux/tag_pipeline.py +++ b/ranboorux/tag_pipeline.py @@ -2,7 +2,8 @@ import re import unicodedata -from typing import Dict, Iterable, List, Optional, Set, Tuple +from dataclasses import dataclass +from typing import Dict, Iterable, List, Optional, Protocol, Set, Tuple, Union # --- Regex Patterns --- _DASH_UNDERSCORE_RE = re.compile(r"[_\-]+") @@ -350,6 +351,23 @@ ) +@dataclass(frozen=True) +class FilterContext: + toggles: Tuple[bool, bool, bool, bool, bool, bool, bool, bool, bool, bool] + base_colors: Tuple[Set[str], Set[str]] + allowed_subjects: Set[str] + cache: Dict[str, str] + favorites_guard: Set[str] + + +class CatalogResolver(Protocol): + def resolve_alias(self, tag: str) -> Optional[str]: ... + def is_textual(self, tag: str) -> bool: ... + def is_hair(self, tag: str) -> bool: ... + def is_eye(self, tag: str) -> bool: ... + def category(self, tag: str) -> int: ... + + # --- Core Tag Pipeline Functions --- @@ -370,6 +388,30 @@ def remove_repeated_tags(prompt: str) -> str: return ",".join(dedupe_keep_order(tags)) +def limit_prompt_tags(prompt: str, limit_val: Union[int, float, str], mode: str) -> str: + tags = split_prompt_tags(prompt) + if not tags: + return "" + if mode == "Limit": + try: + pct = float(limit_val) + except (ValueError, TypeError): + return prompt + if pct <= 0: + return "" + max_count = max(1, int(len(tags) * pct)) + return ",".join(tags[:max_count]) + if mode == "Max": + try: + max_count = int(limit_val) + except (ValueError, TypeError): + return prompt + if max_count <= 0: + return prompt + return ",".join(tags[:max_count]) + return prompt + + def canonicalize_raw_tag(tag: str) -> str: if not isinstance(tag, str): return "" diff --git a/ranboorux/user_store.py b/ranboorux/user_store.py index de5538b..33521ac 100644 --- a/ranboorux/user_store.py +++ b/ranboorux/user_store.py @@ -170,8 +170,9 @@ def append_text_log(file_path: PathLike, lines: Iterable[str]) -> None: target.parent.mkdir(parents=True, exist_ok=True) with target.open("a", encoding="utf-8", newline="\n") as handle: for line in lines: - handle.write(str(line)) - if not str(line).endswith("\n"): + text = str(line) + handle.write(text) + if not text.endswith("\n"): handle.write("\n") except OSError as exc: raise UserStoreError(f"Could not append text log {target}: {exc}") from exc diff --git a/scripts/ranbooru.py b/scripts/ranbooru.py index bffd3ad..b13154d 100644 --- a/scripts/ranbooru.py +++ b/scripts/ranbooru.py @@ -7,21 +7,16 @@ import re import shutil import sys -import time import traceback -import types import unicodedata -import xml.etree.ElementTree as ET from contextlib import ExitStack, contextmanager from datetime import datetime from io import BytesIO from typing import Dict, Iterable, List, Optional, Set, Tuple -from urllib.parse import quote_plus import gradio as gr import modules.scripts as scripts import numpy as np -import requests from modules import shared from modules.processing import ( StableDiffusionProcessing, @@ -37,15 +32,17 @@ from modules.scripts import basedir from ranboorux import catalog as rb_catalog -from ranboorux import host_state as rb_host_state +from ranboorux import http_client as rb_http_client from ranboorux import image_ops as rb_image_ops from ranboorux import loranado as rb_loranado -from ranboorux import prompting as rb_prompting -from ranboorux import requesting as rb_requesting +from ranboorux import mutation_scope as rb_mutation_scope from ranboorux import run_options as rb_run_options from ranboorux import tag_pipeline as rb_tag_pipeline from ranboorux import user_store as rb_user_store +from ranboorux.anima_detect import get_anima_model_info +from ranboorux.boorus import Booru from ranboorux.integrations import adetailer as rb_adetailer_integration +from ranboorux.integrations import adetailer_orchestration as rb_adetailer_orch from ranboorux.integrations import adetailer_runtime as rb_adetailer_runtime from ranboorux.integrations import controlnet as rb_controlnet_integration from ranboorux.integrations import img2img_lifecycle as rb_img2img_lifecycle @@ -134,156 +131,6 @@ _ranbooru_logger = logging.getLogger("ranboorux") -FURRY_CORE_TAGS = { - "anthro", - "furry", - "feral", - "feral_focus", - "feral_only", - "scalie", - "avian", - "hooved_animal", - "digitigrade", - "taur", - "mythological_creature", - "kemono", - "beastman", - "beastgirl", - "beastboy", - "kemonomimi", - "fur", - "fur_focus", -} - -POKEMON_PREFIXES = ( - "pokemon", - "pikachu", - "eevee", - "charizard", - "mewtwo", - "gardevoir", - "lucario", - "lopunny", -) -ANIMAL_EAR_KEYWORDS = ( - "_ear", - "animal_ears", - "beast_ears", - "cat_ears", - "dog_ears", - "fox_ears", - "bunny_ears", - "wolf_ears", - "horse_ears", - "bear_ears", -) -HORN_KEYWORDS = ( - "horn", - "horns", - "antlers", - "unicorn_horn", - "goat_horns", - "demon_horns", - "ram_horns", - "bull_horns", - "long_horns", -) - -HEADWEAR_TAGS = { - "hat", - "cap", - "beret", - "helmet", - "hood", - "crown", - "tiara", - "headband", - "hairband", - "headdress", - "veil", - "witch_hat", - "wizard_hat", - "top_hat", - "beanie", - "goggles", - "glasses_on_head", - "sailor_hat", - "nurse_cap", - "maid_headdress", - "pirate_hat", - "sombrero", - "bunny_ears_headband", - "cat_ears_headband", - "animal_ears_headband", - "motorcycle_helmet", - "baseball_cap", - "bowler_hat", - "straw_hat", - "sun_hat", - "halo", - "circular_halo", - "floating_halo", -} - -HALO_TAGS = {"halo", "circular_halo", "ring_halo", "floating_halo", "angelic_halo"} - -HAIR_COLOR_TAGS = { - "blonde_hair", - "brown_hair", - "black_hair", - "grey_hair", - "gray_hair", - "white_hair", - "silver_hair", - "blue_hair", - "green_hair", - "red_hair", - "pink_hair", - "purple_hair", - "orange_hair", - "aqua_hair", - "magenta_hair", - "teal_hair", - "multicolored_hair", - "gradient_hair", - "rainbow_hair", -} - -EYE_COLOR_TAGS = { - "blue_eyes", - "green_eyes", - "red_eyes", - "brown_eyes", - "black_eyes", - "yellow_eyes", - "amber_eyes", - "orange_eyes", - "purple_eyes", - "pink_eyes", - "golden_eyes", - "silver_eyes", - "grey_eyes", - "gray_eyes", - "white_eyes", - "aqua_eyes", - "heterochromia", - "multicolored_eyes", - "gradient_eyes", -} - - -SERIES_KEYWORDS = { - "franchise", - "series", - "canon", - "official_media", - "gacha_game", - "anime", - "manga_franchise", - "visual_novel", -} -SERIES_SUFFIXES = ("_series", "_franchise", "_media", "_universe") - RATING_TYPES = { "none": {"All": "All"}, "full": {"All": "All", "Safe": "safe", "Questionable": "questionable", "Explicit": "explicit"}, @@ -314,16 +161,6 @@ re.compile(r"(? List[str]: - return [tag.strip() for tag in prompt.split(",") if tag.strip()] - - -def _dedupe_keep_order(tags: Iterable[str]) -> List[str]: - return list(dict.fromkeys(tags)) - - def _split_tag_string(value: Optional[str]) -> List[str]: if not isinstance(value, str): return [] @@ -469,25 +290,6 @@ def _split_tag_string_override(value: object) -> Optional[List[str]]: return _split_tag_string(value) -def remove_repeated_tags(prompt): - try: - return rb_prompting.remove_repeated_tags(prompt) - except Exception as e: - _log(f"Error remove_repeated: {e}. Input: '{prompt}'") - return "" - - -def limit_prompt_tags(prompt, limit_val, mode): - try: - return rb_prompting.limit_prompt_tags(prompt, limit_val, mode) - except ValueError: - _log(f"Error limiting tags: Invalid limit value '{limit_val}'") - return prompt - except Exception as e: - _log(f"Error limiting tags: {e}") - return prompt - - POST_URL_TEMPLATES = { "danbooru": "https://danbooru.donmai.us/posts/{pid}", "gelbooru": "https://gelbooru.com/index.php?page=post&s=view&id={pid}", @@ -521,8 +323,8 @@ def get_original_post_url(post): def generate_chaos(pos_tags, neg_tags, chaos_amount): - pos_tag_list = _split_prompt_tags(pos_tags) - neg_tag_list = _split_prompt_tags(neg_tags) + pos_tag_list = rb_tag_pipeline.split_prompt_tags(pos_tags) + neg_tag_list = rb_tag_pipeline.split_prompt_tags(neg_tags) chaos_list = list(set(pos_tag_list + neg_tag_list)) if not chaos_list: return pos_tags, neg_tags @@ -532,639 +334,15 @@ def generate_chaos(pos_tags, neg_tags, chaos_amount): pos_add = chaos_list[len_list:] final_pos = list(set(pos_tag_list) - set(neg_add)) + pos_add final_neg = list(set(neg_tag_list) - set(pos_add)) + neg_add - return ",".join(_dedupe_keep_order(final_pos)), ",".join(_dedupe_keep_order(final_neg)) + return ",".join(rb_tag_pipeline.dedupe_keep_order(final_pos)), ",".join( + rb_tag_pipeline.dedupe_keep_order(final_neg) + ) class BooruError(Exception): pass -class Booru: - def __init__(self, booru_name, base_api_url, http_client=None): - self.booru_name = booru_name - self.base_api_url = base_api_url - self.http = http_client or rb_requesting.BooruSession() - self.headers = {"user-agent": f"Ranbooru Extension/{Script.version} for Forge"} - - def _fetch_data(self, query_url): - _log(f"Querying {self.booru_name}: {rb_requesting.redact_url(query_url)}") - try: - return self.http.get_json(query_url, headers=self.headers, timeout=30) - except Exception as e: - status_code = getattr(getattr(e, "response", None), "status_code", None) - if status_code == 401 or "401" in str(e): - err_msg = ( - f"Authentication failed (401 Unauthorized) for {self.booru_name}. " - "Please check your API key and User ID under Gelbooru settings. " - "Note: User ID must be your numeric account ID (e.g. 123456), not your username." - ) - _log(f"Error {err_msg}") - raise BooruError(err_msg) from e - message = rb_requesting.safe_exception_message( - f"fetching data from {self.booru_name}", query_url, e - ) - _log(f"Error {message}") - raise BooruError(f"HTTP Error {message}") from e - - def _is_direct_image_url(self, url): - """Check if URL is a direct image URL (not from external sites like Pixiv/Twitter)""" - if not url or not isinstance(url, str): - return False - - # Skip external sites that don't provide direct image access - external_sites = [ - "pixiv.net", - "pximg.net", - "twitter.com", - "x.com", - "t.co", - "deviantart.com", - "artstation.com", - "instagram.com", - "facebook.com", - "patreon.com", - "fanbox.cc", - ] - - url_lower = url.lower() - for site in external_sites: - if site in url_lower: - return False - - # Check if URL ends with common image extensions - image_extensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff"] - if any(url_lower.endswith(ext) for ext in image_extensions): - return True - - # Check if URL contains image-serving patterns - if any(pattern in url_lower for pattern in ["/images/", "/img/", "/media/", "/files/"]): - return True - - return False - - def _standardize_post(self, post_data): - post = {} - # extract tags in a robust way; some APIs return categorized tags as dicts - raw_tags = post_data.get("tags", post_data.get("tag_string", "")) - # store categorized lists when possible - artist_tags = [] - character_tags = [] - copyright_tags = [] - if isinstance(post_data.get("tags"), dict): - tags_dict = post_data.get("tags") - # e621 style: tags dict with sublevels - if isinstance(tags_dict.get("artist"), list): - artist_tags = tags_dict.get("artist", []) - if isinstance(tags_dict.get("character"), list): - character_tags = tags_dict.get("character", []) - if isinstance(tags_dict.get("copyright"), list): - copyright_tags = tags_dict.get("copyright", []) - if "tag_string_artist" in post_data: - parsed = _split_tag_string_override(post_data.get("tag_string_artist")) - if parsed is not None: - artist_tags = parsed - if "tag_string_character" in post_data: - parsed = _split_tag_string_override(post_data.get("tag_string_character")) - if parsed is not None: - character_tags = parsed - if "tag_string_copyright" in post_data: - parsed = _split_tag_string_override(post_data.get("tag_string_copyright")) - if parsed is not None: - copyright_tags = parsed - - # For boorus that don't provide categorized tags, try to extract character tags from the main tag string - # This handles cases like Gelbooru/Danbooru where character tags are mixed with other tags - if not character_tags and isinstance(raw_tags, str): - all_tags = _split_tag_string(raw_tags) - for tag in all_tags: - # Common patterns for character tags: contains parentheses (series name) or ends with specific patterns - if ( - ("(" in tag and ")" in tag) - or tag.endswith(r"_\(series\)") - or tag.endswith(r"_\(character\)") - ): - character_tags.append(tag) - # Also catch some common character name patterns (this is heuristic but should catch most) - elif any( - series in tag.lower() - for series in [ - "genshin_impact", - "touhou", - "fate_", - "azur_lane", - "kantai_collection", - "pokemon", - ] - ): - character_tags.append(tag) - - post["tags"] = raw_tags - post["artist_tags"] = artist_tags - post["character_tags"] = character_tags - post["copyright_tags"] = copyright_tags - post["score"] = post_data.get("score", 0) - post["file_url"] = post_data.get("file_url") - if post["file_url"] is None: - post["file_url"] = post_data.get("large_file_url") - if post["file_url"] is None: - # Check if source is a direct image URL before using it - source_url = post_data.get("source") - if source_url and self._is_direct_image_url(source_url): - post["file_url"] = source_url - else: - post["file_url"] = None - post["id"] = post_data.get("id") - post["rating"] = post_data.get("rating") - post["booru_name"] = self.booru_name - return post - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - raise NotImplementedError - - -class Gelbooru(Booru): - def __init__(self, fringe_benefits, credentials: Optional[Dict[str, str]] = None): - super().__init__( - "Gelbooru", - f"https://gelbooru.com/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", - ) - self.fringeBenefits = fringe_benefits - credentials = credentials or {} - self.api_key = ( - _sanitize_gelbooru_credential(credentials.get("api_key")) - if isinstance(credentials, dict) - else "" - ) - self.user_id = ( - _sanitize_gelbooru_credential(credentials.get("user_id")) - if isinstance(credentials, dict) - else "" - ) - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if not self.api_key or not self.user_id: - raise BooruError( - "Gelbooru requires an API key and user ID. Set them under RanbooruX ? Gelbooru settings." - ) - credentials_query = ( - f"&api_key={quote_plus(self.api_key)}&user_id={quote_plus(self.user_id)}" - ) - if post_id: - query_url = f"{self.base_api_url}{credentials_query}&id={post_id}{tags_query}" - fetched_data = self._fetch_data(query_url) - if fetched_data and "post" in fetched_data and isinstance(fetched_data["post"], list): - all_fetched_posts = fetched_data["post"] - COUNT = len(all_fetched_posts) - print(f"[R] Found {COUNT} post(s) for ID: {post_id}") - else: - page = random.randint(0, max_pages - 1) - query_url = f"{self.base_api_url}{credentials_query}&pid={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if fetched_data and "post" in fetched_data and isinstance(fetched_data["post"], list): - all_fetched_posts = fetched_data["post"] - if ( - fetched_data - and "@attributes" in fetched_data - and "count" in fetched_data["@attributes"] - ): - try: - COUNT = int(fetched_data["@attributes"]["count"]) - except Exception: - COUNT = len(all_fetched_posts) - else: - COUNT = len(all_fetched_posts) - print( - f"[R] Fetched {len(all_fetched_posts)} posts from page {page}. Reported total (approx): {COUNT}" - ) - return [self._standardize_post(post) for post in all_fetched_posts] - - -class GelbooruCompatible(Booru): - RETRIABLE_STATUS = {429, 500, 502, 503, 504} - - def __init__( - self, base_url: str, retries: int = 3, backoff: float = 1.5, log_diagnostics: bool = True - ): - sanitized = _sanitize_gelbooru_compat_base_url(base_url) - if not sanitized: - raise ValueError("Invalid Gelbooru-compatible base URL.") - self.base_url = sanitized - self.retries = max(1, retries) - self.backoff = max(0.5, backoff) - self.log_diagnostics = log_diagnostics - self._post_endpoint = f"{self.base_url}/index.php?page=dapi&s=post&q=index" - self._tag_endpoint = f"{self.base_url}/index.php?page=dapi&s=tag&q=index" - self._alias_endpoint = f"{self.base_url}/index.php?page=dapi&s=tag_alias&q=index" - super().__init__("Gelbooru-Compatible", self._post_endpoint) - - def _perform_request(self, url: str) -> requests.Response: - last_error: Optional[Exception] = None - for attempt in range(1, self.retries + 1): - try: - response = self.http.get(url, headers=self.headers, timeout=30, stream=True) - except Exception as exc: - last_error = exc - self._log_retry(url, attempt, f"Request error: {exc.__class__.__name__}") - else: - if response.status_code in self.RETRIABLE_STATUS: - last_error = BooruError(f"Status {response.status_code}") - self._log_retry(url, attempt, f"Status {response.status_code}") - close = getattr(response, "close", None) - if callable(close): - close() - else: - content = self.http._read_bounded_response( - response, - url, - rb_requesting.DEFAULT_API_MAX_BYTES, - ) - return rb_requesting.BoundedResponse( - url=str(getattr(response, "url", url) or url), - status_code=int(getattr(response, "status_code", 200) or 200), - headers=getattr(response, "headers", {}) or {}, - content=content, - encoding=getattr(response, "encoding", None), - ) - time.sleep(min(self.backoff * attempt, 5.0)) - if last_error is None: - error_summary = "unknown error" - elif isinstance(last_error, BooruError): - error_summary = str(last_error) - else: - error_summary = last_error.__class__.__name__ - raise BooruError( - f"HTTP Error fetching from {self.booru_name}: {error_summary} for {rb_requesting.redact_url(url)}" - ) - - def _log_retry(self, url: str, attempt: int, message: str) -> None: - _log(f"{self.booru_name}: retry {attempt} for {rb_requesting.redact_url(url)} - {message}") - - def _log_snippet(self, response: requests.Response) -> None: - if not self.log_diagnostics: - return - snippet = response.text.strip().replace("\n", " ")[:200] - _log( - f"{self.booru_name}: {rb_requesting.redact_url(getattr(response, 'url', ''))} -> {snippet}" - ) - - def _parse_json_entities(self, payload, entity_key: str) -> Tuple[List[dict], Optional[int]]: - entries: List[dict] = [] - approx = None - if isinstance(payload, dict): - possible = payload.get(entity_key) - if isinstance(possible, list): - entries = possible - elif isinstance(possible, dict): - entries = [possible] - attrs = payload.get("@attributes") - if isinstance(attrs, dict) and "count" in attrs: - try: - approx = int(attrs["count"]) - except (TypeError, ValueError): - approx = None - elif isinstance(payload, list): - entries = payload - return entries, approx - - def _parse_xml_entities( - self, text_payload: str, entity_key: str - ) -> Tuple[List[dict], Optional[int]]: - probe = (text_payload or "").lower() - if (" Tuple[List[dict], int]: - json_url = f"{url_base}&json=1" - try: - response = self._perform_request(json_url) - self._log_snippet(response) - ct = (response.headers.get("content-type") or "").lower() - text_head = (response.text or "").lstrip()[:64].lower() - if ( - "html" in ct - or text_head.startswith(" 0 else 0 - query_base = f"{self._post_endpoint}&limit={POST_AMOUNT}&pid={page}{tags_query}" - posts, approx = self._request_dapi(query_base, "post") - COUNT = approx - print( - f"[R] Gelbooru-compatible: fetched {len(posts)} posts from page {page}. Reported count={approx}" - ) - standardized = [] - for post in posts: - normalized = self._standardize_post(post) - normalized["source_base_url"] = self.base_url - standardized.append(normalized) - return standardized - - def get_tags(self, name_pattern: Optional[str] = None, limit: int = 100) -> List[dict]: - query = f"{self._tag_endpoint}&limit={limit}" - if name_pattern: - query += f"&name_pattern={quote_plus(name_pattern)}" - tags, _ = self._request_dapi(query, "tag") - return tags - - def get_tag_aliases(self, name_pattern: Optional[str] = None, limit: int = 100) -> List[dict]: - query = f"{self._alias_endpoint}&limit={limit}" - if name_pattern: - query += f"&name_pattern={quote_plus(name_pattern)}" - aliases, _ = self._request_dapi(query, "tag_alias") - return aliases - - -class Danbooru(Booru): - def __init__(self): - super().__init__("Danbooru", f"https://danbooru.donmai.us/posts.json?limit={POST_AMOUNT}") - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - query_url = f"https://danbooru.donmai.us/posts/{post_id}.json" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, dict) and "id" in fetched_data: - all_fetched_posts = [fetched_data] - COUNT = len(all_fetched_posts) - print(f"[R] Found {COUNT} post(s) for ID: {post_id}") - else: - page = random.randint(1, max_pages) - query_url = f"{self.base_api_url}&page={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, list): - all_fetched_posts = fetched_data - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from page {page}.") - return [self._standardize_post(post) for post in all_fetched_posts if post] - - -class XBooru(Booru): - def __init__(self): - super().__init__( - "XBooru", - f"https://xbooru.com/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", - ) - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - query_url = f"{self.base_api_url}&id={post_id}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, dict) and "id" in fetched_data: - all_fetched_posts = [fetched_data] - else: - page = random.randint(0, max_pages - 1) - query_url = f"{self.base_api_url}&pid={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, list): - all_fetched_posts = fetched_data - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from XBooru.") - standardized_posts = [] - for post_data in all_fetched_posts: - post = self._standardize_post(post_data) - if "directory" in post_data and "image" in post_data: - post["file_url"] = ( - f"https://xbooru.com/images/{post_data['directory']}/{post_data['image']}" - ) - standardized_posts.append(post) - return standardized_posts - - -class Rule34(Booru): - def __init__(self): - super().__init__( - "Rule34", - f"https://api.rule34.xxx/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", - ) - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - query_url = f"{self.base_api_url}&id={post_id}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, dict) and "id" in fetched_data: - all_fetched_posts = [fetched_data] - else: - page = random.randint(0, max_pages - 1) - query_url = f"{self.base_api_url}&pid={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, list): - all_fetched_posts = fetched_data - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from Rule34.") - return [self._standardize_post(post) for post in all_fetched_posts] - - -class Safebooru(Booru): - def __init__(self): - super().__init__( - "Safebooru", - f"https://safebooru.org/index.php?page=dapi&s=post&q=index&json=1&limit={POST_AMOUNT}", - ) - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - query_url = f"{self.base_api_url}&id={post_id}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, dict) and "id" in fetched_data: - all_fetched_posts = [fetched_data] - else: - page = random.randint(0, max_pages - 1) - query_url = f"{self.base_api_url}&pid={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, list): - all_fetched_posts = fetched_data - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from Safebooru.") - standardized_posts = [] - for post_data in all_fetched_posts: - post = self._standardize_post(post_data) - if "directory" in post_data and "image" in post_data: - post["file_url"] = ( - f"https://safebooru.org/images/{post_data['directory']}/{post_data['image']}" - ) - standardized_posts.append(post) - return standardized_posts - - -class Konachan(Booru): - def __init__(self): - super().__init__("Konachan", f"https://konachan.com/post.json?limit={POST_AMOUNT}") - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - print("[R] Warn: Konachan does not support post IDs.") - return [] - page = random.randint(1, max_pages) - query_url = f"{self.base_api_url}&page={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, list): - all_fetched_posts = fetched_data - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from Konachan.") - return [self._standardize_post(post) for post in all_fetched_posts] - - -class Yandere(Booru): - def __init__(self): - super().__init__("Yandere", f"https://yande.re/post.json?limit={POST_AMOUNT}") - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - print("[R] Warn: Yandere does not support post IDs.") - return [] - page = random.randint(1, max_pages) - query_url = f"{self.base_api_url}&page={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, list): - all_fetched_posts = fetched_data - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from Yandere.") - return [self._standardize_post(post) for post in all_fetched_posts] - - -class AIBooru(Booru): - def __init__(self): - super().__init__("AIBooru", f"https://aibooru.online/posts.json?limit={POST_AMOUNT}") - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - print("[R] Warn: AIBooru does not support post IDs.") - return [] - page = random.randint(1, max_pages) - query_url = f"{self.base_api_url}?limit={POST_AMOUNT}&page={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if isinstance(fetched_data, list): - all_fetched_posts = fetched_data - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from AIBooru.") - standardized_posts = [] - for post_data in all_fetched_posts: - post = self._standardize_post(post_data) - post["tags"] = post_data.get("tag_string", "") - standardized_posts.append(post) - return standardized_posts - - -class e621(Booru): - def __init__(self): - super().__init__("e621", f"https://e621.net/posts.json?limit={POST_AMOUNT}") - - def get_posts(self, tags_query="", max_pages=10, post_id=None): - global COUNT - COUNT = 0 - all_fetched_posts = [] - if post_id: - print("[R] Warn: e621 does not support post IDs.") - return [] - page = random.randint(1, max_pages) - query_url = f"{self.base_api_url}?page={page}{tags_query}" - fetched_data = self._fetch_data(query_url) - if ( - isinstance(fetched_data, dict) - and "posts" in fetched_data - and isinstance(fetched_data["posts"], list) - ): - all_fetched_posts = fetched_data["posts"] - COUNT = len(all_fetched_posts) - print(f"[R] Fetched {COUNT} posts from e621.") - standardized_posts = [] - for post_data in all_fetched_posts: - post = self._standardize_post(post_data) - temp_tags = [] - sublevels = ["general", "artist", "copyright", "character", "species"] - if "tags" in post_data: - for sublevel in sublevels: - if sublevel in post_data["tags"] and isinstance( - post_data["tags"][sublevel], list - ): - temp_tags.extend(post_data["tags"][sublevel]) - post["tags"] = " ".join(temp_tags) - if ( - "score" in post_data - and isinstance(post_data["score"], dict) - and "total" in post_data["score"] - ): - post["score"] = post_data["score"]["total"] - standardized_posts.append(post) - return standardized_posts - - class TagCatalogProvider: """Interface for optional tag catalog backends.""" @@ -1449,6 +627,7 @@ def __init__(self): self._gelbooru_effective_credentials: Optional[Dict[str, str]] = None self._personal_remove_tags: Set[str] = set() self._favorite_tags: Set[str] = set() + self._is_anima_model: bool = False self._removal_context: Dict[str, object] = {} self._tag_normal_cache: Dict[str, str] = {} self._synonym_groups: Tuple[Set[str], ...] = tuple() @@ -1471,7 +650,8 @@ def __init__(self): self._synonym_lookup = {} self._adetailer_state = rb_adetailer_runtime.AdetailerRunState() self._adetailer_patches = rb_adetailer_runtime.PatchRegistry() - self._host_scope = rb_host_state.HostMutationScope() + self._host_scope = rb_mutation_scope.HostMutationScope() + self._adetailer_orch = rb_adetailer_orch.AdetailerOrchestrator(self) self._strict_img2img_fetch: bool = True self._strict_img2img_active: bool = False self._strict_img2img_relaxed: bool = False @@ -1490,7 +670,7 @@ def __init__(self): self._tag_catalog_linter_limit: int = 3 self._catalog_subject_anchors = None self._loranado_scan_cache: Dict[str, Dict[str, object]] = {} - self._http_client = rb_requesting.BooruSession(use_cache=False) + self._http_client = rb_http_client.BooruSession(use_cache=False) self._load_tag_catalog_preferences() sorting_priority = 1 # Highest priority to run before ALL other extensions @@ -1512,187 +692,6 @@ def __init__(self): "personal": PERSONAL_REMOVE_FILE, "favorites": FAVORITES_FILE, } - _CLOTHING_KEYWORDS = { - "dress", - "shirt", - "skirt", - "skorts", - "pants", - "jeans", - "shorts", - "jacket", - "coat", - "sweater", - "hoodie", - "kimono", - "robe", - "uniform", - "school uniform", - "sailor uniform", - "bikini", - "swimsuit", - "lingerie", - "underwear", - "panties", - "bra", - "corset", - "thighhighs", - "stockings", - "socks", - "gloves", - "mittens", - "scarf", - "cape", - "apron", - "armor", - "bustier", - "bodysuit", - "leotard", - "gown", - "tuxedo", - "suit", - "vest", - "necktie", - "bowtie", - "hat", - "cap", - "headband", - "hairband", - "headdress", - "veil", - "crown", - "helmet", - "sandals", - "boots", - "shoes", - "heels", - "sneakers", - "flip flops", - "garter", - "garter belt", - "pantyhose", - "stocking", - "cloak", - "cardigan", - "sleeves", - "armband", - "choker", - "ribbon", - "bow", - "shawl", - "loincloth", - "loin cloth", - "tabard", - "capelet", - "poncho", - "overalls", - "tank top", - "t-shirt", - "tee shirt", - "pajamas", - "nightgown", - } - _TEXTUAL_TAGS = { - "text", - "english text", - "japanese text", - "chinese text", - "korean text", - "translated", - "translation", - "commentary", - "artist commentary", - "author commentary", - "publisher commentary", - "copyright text", - "speech bubble", - "speech bubbles", - "dialogue", - "dialog", - "sound effect", - "sound effects", - "comic text", - "comic panel", - "subtitle", - "subtitles", - "caption", - "captions", - "floating text", - "text focus", - "text overlay", - "text background", - "watermark", - "watermark text", - "signature", - "sign", - "tagme", - "written text", - "scribble", - "handwritten text", - "handwriting", - "text box", - "thought bubble", - "thought balloon", - "logo", - "logo text", - "notice", - "speech bubble text", - } - _SUBJECT_TAGS = { - "solo", - "duo", - "trio", - "quartet", - "group", - "gang", - "crowd", - "couple", - "threesome", - "foursome", - "orgy", - "1girl", - "2girls", - "3girls", - "4girls", - "1boy", - "2boys", - "3boys", - "4boys", - "1other", - "2others", - "3others", - "4others", - "multiple girls", - "multiple boys", - "multiple people", - "multiple others", - "solo focus", - "female focus", - "male focus", - "mixed group", - "1female", - "1male", - "2females", - "2males", - "3females", - "3males", - "1person", - "2people", - "3people", - "4people", - } - _FURRY_CORE_NORMALIZED = {tag.replace("_", " ") for tag in FURRY_CORE_TAGS} - _POKEMON_PREFIXES_NORMALIZED = tuple(prefix.replace("_", " ") for prefix in POKEMON_PREFIXES) - _ANIMAL_EAR_KEYWORDS_NORMALIZED = tuple( - keyword.replace("_", " ") for keyword in ANIMAL_EAR_KEYWORDS - ) - _HORN_KEYWORDS_NORMALIZED = tuple(keyword.replace("_", " ") for keyword in HORN_KEYWORDS) - _HEADWEAR_TAGS_NORMALIZED = {tag.replace("_", " ") for tag in HEADWEAR_TAGS} - _HALO_TAGS_NORMALIZED = {tag.replace("_", " ") for tag in HALO_TAGS} - _HAIR_COLOR_TAGS_NORMALIZED = {tag.replace("_", " ") for tag in HAIR_COLOR_TAGS} - _EYE_COLOR_TAGS_NORMALIZED = {tag.replace("_", " ") for tag in EYE_COLOR_TAGS} - _SERIES_KEYWORDS_NORMALIZED = {tag.replace("_", " ") for tag in SERIES_KEYWORDS} - _SERIES_SUFFIXES_NORMALIZED = tuple(suffix.replace("_", " ") for suffix in SERIES_SUFFIXES) @staticmethod def _canonicalize_raw_tag(tag: str) -> str: @@ -1969,7 +968,7 @@ def _apply_optional_catalog( subject_anchors = getattr(self, "_catalog_subject_anchors", None) if subject_anchors is None: - subject_anchors = {s.replace(" ", "_") for s in getattr(self, "_SUBJECT_TAGS", set())} + subject_anchors = {s.replace(" ", "_") for s in rb_tag_pipeline._SUBJECT_TAGS} self._catalog_subject_anchors = subject_anchors kept: List[str] = [] @@ -2400,48 +1399,6 @@ def _ui_set_gelbooru_compat_base_url(self, base_url: Optional[str]): self._gelbooru_compat_base_url = sanitized return _gr_component_update(gr.Textbox, value=self._gelbooru_compat_base_url) - def _is_furry_tag(self, tag: str) -> bool: - return rb_tag_pipeline.is_furry_tag(tag) - - def _is_headwear_tag(self, tag: str) -> bool: - return rb_tag_pipeline.is_headwear_tag(tag) - - def _is_girl_suffix_tag(self, tag: str) -> bool: - return rb_tag_pipeline.is_girl_suffix_tag(tag) - - def _is_hair_color_tag(self, tag: str) -> bool: - normalized = (self._normalize_tag(tag) or "").strip().lower() - if not normalized: - normalized = self._canonicalize_raw_tag(tag) - if not normalized: - return False - catalog = self._active_catalog() - if catalog and catalog.is_hair(normalized.replace(" ", "_")): - return True - return rb_tag_pipeline.is_hair_color_tag(tag) - - def _is_eye_color_tag(self, tag: str) -> bool: - normalized = (self._normalize_tag(tag) or "").strip().lower() - if not normalized: - normalized = self._canonicalize_raw_tag(tag) - if not normalized: - return False - catalog = self._active_catalog() - if catalog and catalog.is_eye(normalized.replace(" ", "_")): - return True - return rb_tag_pipeline.is_eye_color_tag(tag) - - def _is_series_tag(self, tag: str) -> bool: - normalized = (self._normalize_tag(tag) or "").strip().lower() - if not normalized: - normalized = self._canonicalize_raw_tag(tag) - if not normalized: - return False - catalog = self._active_catalog() - if catalog and catalog.category(normalized.replace(" ", "_")) == 3: - return True - return rb_tag_pipeline.is_series_tag(tag) - def _extract_color_tags(self, text: str) -> tuple[set[str], set[str]]: hair_tags: set[str] = set() eye_tags: set[str] = set() @@ -2463,27 +1420,12 @@ def _extract_color_tags(self, text: str) -> tuple[set[str], set[str]]: if catalog.is_eye(token_key): canonical = catalog.resolve_alias(token_key) eye_tags.add(canonical.replace("_", " ") if canonical else normalized) - if normalized in self._HAIR_COLOR_TAGS_NORMALIZED: + if normalized in rb_tag_pipeline._HAIR_COLOR_TAGS_NORMALIZED: hair_tags.add(normalized) - if normalized in self._EYE_COLOR_TAGS_NORMALIZED: + if normalized in rb_tag_pipeline._EYE_COLOR_TAGS_NORMALIZED: eye_tags.add(normalized) return hair_tags, eye_tags - def _is_clothing_tag(self, tag: str) -> bool: - return rb_tag_pipeline.is_clothing_tag(tag) - - def _is_textual_tag(self, tag: str) -> bool: - normalized = self._normalize_tag(tag) - if not normalized: - return False - catalog = self._active_catalog() - if catalog and catalog.is_textual(normalized.replace(" ", "_")): - return True - return rb_tag_pipeline.is_textual_tag(tag) - - def _is_subject_tag(self, tag: str) -> bool: - return rb_tag_pipeline.is_subject_tag(tag) - def _extract_subject_tags(self, text: str) -> set: return rb_tag_pipeline.extract_subject_tags(text) @@ -2819,363 +1761,272 @@ def refresh_ser(self): def refresh_rem(self): return _gr_update(choices=self.get_files(USER_REMOVE_DIR)) - def ui(self, is_img2img): - with InputAccordion(False, label="RanbooruX", elem_id=self.elem_id("ra_enable")) as enabled: - booru_list = [ - "danbooru", - "gelbooru", - "gelbooru-compatible", - "xbooru", - "rule34", - "safebooru", - "konachan", - "yande.re", - "aibooru", - "e621", - ] - booru = gr.Dropdown(booru_list, label="Booru", value="danbooru") - with gr.Group(visible=False) as gelbooru_credentials_group: - gelbooru_api_key = gr.Textbox( - label="Gelbooru API Key", - type="password", - placeholder="Enter your Gelbooru API key", - ) - gelbooru_user_id = gr.Textbox( - label="Gelbooru User ID", - placeholder="Numeric account ID e.g. 123456 (not username)", - info="Numeric User ID from Gelbooru Options page (e.g. 123456, NOT your username)", - ) - gelbooru_save_button = gr.Button("Save Credentials to Disk", variant="primary") - gelbooru_saved_message = gr.Markdown("", visible=False) - gelbooru_clear_button = gr.Button("Clear Saved Credentials", visible=False) - with gr.Group(visible=False) as gelbooru_compat_group: - gelbooru_compat_base_url = gr.Textbox( - label="Gelbooru-compatible Base URL", - placeholder="https://realbooru.com", - value=self._gelbooru_compat_base_url, - ) - max_pages = gr.Slider( - label="Max Pages (tag search)", minimum=1, maximum=100, value=10, step=1 - ) - gr.Markdown("""## Post""") - post_id = gr.Textbox(lines=1, label="Post ID (Overrides tags/pages)") - gr.Markdown("""## Tags""") - tags = gr.Textbox(lines=1, label="Tags to Search (Pre)") - remove_tags = gr.Textbox(lines=1, label="Tags to Remove (Post)") - mature_rating = gr.Radio( - list(RATINGS.get("gelbooru", RATING_TYPES["none"])), - label="Mature Rating", - value="All", - ) - with gr.Accordion("Removal Filters", open=False): - with gr.Group(): - gr.Markdown("**Danbooru Tag Catalog**") + def _build_catalog_ui_section(self): + """Tag catalog toggle, file picker, validation/import, and diagnostics. - use_tag_catalog = gr.Checkbox( - label="Use Danbooru Tag Catalog", - value=bool(self._use_tag_catalog), - info="Enable category-aware filtering and alias resolution.", - ) + Creates the Danbooru Tag Catalog group (toggle, source, custom path, import, + validation, reload, status) and the Platform Diagnostics toggle. Returns the + two components that must appear in the script-args component list. - catalog_source = gr.Radio( - ["Bundled", "Custom file"], - label="Catalog Source", - value=("Custom file" if self._catalog_source == "custom" else "Bundled"), - visible=bool(self._use_tag_catalog), - ) + Must be called inside ``gr.Group()`` that lives inside the Removal Filters + accordion. + """ + gr.Markdown("**Danbooru Tag Catalog**") - with gr.Group( - visible=bool(self._use_tag_catalog and self._catalog_source == "custom") - ) as custom_catalog_group: - catalog_upload = gr.File( - label="Upload CSV", file_types=[".csv"], file_count="single" - ) - catalog_path = gr.Textbox( - label="Custom CSV Path", - value=self._custom_catalog_path, - placeholder="/path/to/custom_catalog.csv", - ) - with gr.Row(): - catalog_import_btn = gr.Button("Import Custom Catalog") - catalog_validate_btn = gr.Button("Validate CSV") + use_tag_catalog = gr.Checkbox( + label="Use Danbooru Tag Catalog", + value=bool(self._use_tag_catalog), + info="Enable category-aware filtering and alias resolution.", + ) - reload_catalog = gr.Button( - "Reload Catalog", visible=bool(self._use_tag_catalog) - ) - catalog_status = gr.Markdown( - self._tag_catalog_status_text or "Catalog mode: OFF" - ) + catalog_source = gr.Radio( + ["Bundled", "Custom file"], + label="Catalog Source", + value=("Custom file" if self._catalog_source == "custom" else "Bundled"), + visible=bool(self._use_tag_catalog), + ) - self._catalog_status_md = catalog_status - self._tag_diag_md = None + with gr.Group( + visible=bool(self._use_tag_catalog and self._catalog_source == "custom") + ) as custom_catalog_group: + catalog_upload = gr.File(label="Upload CSV", file_types=[".csv"], file_count="single") + catalog_path = gr.Textbox( + label="Custom CSV Path", + value=self._custom_catalog_path, + placeholder="/path/to/custom_catalog.csv", + ) + with gr.Row(): + catalog_import_btn = gr.Button("Import Custom Catalog") + catalog_validate_btn = gr.Button("Validate CSV") - gr.Markdown("**Quick Presets**: apply common filter combinations with one click.") + reload_catalog = gr.Button("Reload Catalog", visible=bool(self._use_tag_catalog)) + catalog_status = gr.Markdown(self._tag_catalog_status_text or "Catalog mode: OFF") - with gr.Row(): - preset_strip_series = gr.Button("Strip Series/Character") - preset_remove_text = gr.Button("Remove Text-like Tags") - preset_preserve_colors = gr.Button("Preserve Base Colors") - preset_quick_strip = gr.Button("Quick Strip") - with gr.Group(): - gr.Markdown("**Text & Metadata**") - remove_bad_tags = gr.Checkbox( - label="Remove common 'bad' tags", - value=True, - info="Cull frequent watermark, commentary, and UI text tags from prompts.", - ) - remove_text_tags = gr.Checkbox( - label="Remove tag/text/commentary metadata", - value=True, - info="Strip speech bubbles, watermark text, and similar metadata from fetched prompts.", - ) - with gr.Group(): - gr.Markdown("**Characters & Series**") - remove_artist_tags = gr.Checkbox( - label="Remove artist tags", - value=False, - info="Drop artist credits drawn from the source post.", - ) - remove_character_tags = gr.Checkbox( - label="Remove character tags", - value=False, - info="Filter character/franchise tags sourced from metadata.", - ) - remove_series_tags = gr.Checkbox( - label="Remove series / franchise tags", - value=False, - info="Ignore franchise/game/anime tags to keep prompts generic.", - ) - with gr.Group(): - gr.Markdown("**Clothing & Accessories**") - remove_clothing_tags = gr.Checkbox( - label="Remove clothing tags", - value=False, - info="Omit apparel/accessory tags introduced by the booru.", - ) - with gr.Group(): - gr.Markdown("**Furry & Headwear**") - remove_furry_tags = gr.Checkbox( - label="Filter furry/pokemon tags", - value=False, - info="Remove furry, pokemon, and animal trait tags.", - ) - remove_headwear_tags = gr.Checkbox( - label="Filter headwear / halo tags", - value=False, - info="Strip hats, halos, and similar head accessories.", - ) - with gr.Group(): - gr.Markdown("**Girl Suffix**") - remove_girl_suffix_tags = gr.Checkbox( - label="Filter _girl suffix tags", - value=False, - info="Remove demon_girl, cat_girl, angel_girl and similar *_girl tags (keeps 1girl, 2girls, etc.).", - ) - with gr.Group(): + self._catalog_status_md = catalog_status + self._tag_diag_md = None - gr.Markdown("**Colors & Traits**") - preserve_hair_eye_colors = gr.Checkbox( - label="Preserve base hair & eye colors", - value=False, - info="Keep your prompt's hair/eye colors while removing conflicting imports.", - ) - with gr.Group(): - gr.Markdown("**Subject Constraints**") - restrict_subject_tags = gr.Checkbox( - label="Keep only subject counts", - value=False, - info="Maintain your subject count (e.g., solo/1girl) by removing mismatched tags.", - ) - personal_choices = self._read_list_file(PERSONAL_REMOVE_FILE) - favorite_choices = self._read_list_file(FAVORITES_FILE) - with gr.Accordion("Personal Lists", open=False): - with gr.Row(): - with gr.Column(): - gr.Markdown("**Personal Removal List**") - personal_remove_dropdown = gr.Dropdown( - choices=personal_choices, - value=personal_choices, - multiselect=True, - label="Removal Tags", - allow_custom_value=False, - ) - personal_remove_input = gr.Textbox( - label="Add tags", placeholder="comma or newline separated" - ) - with gr.Row(): - personal_add_btn = gr.Button("Add", variant="primary") - personal_remove_btn = gr.Button("Remove Selected") - personal_dedupe_btn = gr.Button("De-duplicate") - with gr.Row(): - personal_import_file = gr.File( - label="Import CSV/TXT", file_types=[".txt", ".csv"], visible=True - ) - personal_export_btn = gr.DownloadButton("Export") - with gr.Column(): - gr.Markdown("**Favorites List**") - favorites_dropdown = gr.Dropdown( - choices=favorite_choices, - value=favorite_choices, - multiselect=True, - label="Favorite Tags", - allow_custom_value=False, - ) - favorites_input = gr.Textbox( - label="Add favorites", placeholder="comma or newline separated" - ) - with gr.Row(): - favorites_add_btn = gr.Button("Add", variant="primary") - favorites_remove_btn = gr.Button("Remove Selected") - favorites_dedupe_btn = gr.Button("De-duplicate") - with gr.Row(): - favorites_import_file = gr.File( - label="Import CSV/TXT", file_types=[".txt", ".csv"], visible=True - ) - favorites_export_btn = gr.DownloadButton("Export") - shuffle_tags = gr.Checkbox(label="Shuffle tags", value=True) - change_dash = gr.Checkbox(label='Convert "_" to spaces', value=False) - same_prompt = gr.Checkbox(label="Use same prompt for batch", value=False) - fringe_benefits = gr.Checkbox( - label="Gelbooru: Fringe Benefits", value=True, visible=False - ) - limit_tags = gr.Slider( - value=1.0, label="Limit tags by %", minimum=0.05, maximum=1.0, step=0.05 - ) - max_tags = gr.Slider( - value=0, label="Max tags (0=disabled)", minimum=0, maximum=300, step=1 - ) - change_background = gr.Radio( - ["Don't Change", "Add Detail", "Force Simple", "Force Transparent/White"], - label="Change Background", - value="Don't Change", - ) - change_color = gr.Radio( - ["Don't Change", "Force Color", "Force Monochrome"], - label="Change Color", - value="Don't Change", - ) - sorting_order = gr.Radio( - ["Random", "Score Descending", "Score Ascending"], - label="Sort Order (tag search)", - value="Random", - ) - booru.change(get_available_ratings, booru, mature_rating) - booru.change(show_fringe_benefits, booru, fringe_benefits) - booru.change( - self._update_gelbooru_ui_visibility, - inputs=[booru], - outputs=[ - gelbooru_credentials_group, - gelbooru_saved_message, - gelbooru_clear_button, - gelbooru_api_key, - gelbooru_user_id, - ], - queue=False, - ) - booru.change( - self._update_gelbooru_compat_visibility, - inputs=[booru], - outputs=[gelbooru_compat_group, gelbooru_compat_base_url], - queue=False, - ) - gelbooru_compat_base_url.change( - fn=self._ui_set_gelbooru_compat_base_url, - inputs=[gelbooru_compat_base_url], - outputs=[gelbooru_compat_base_url], - queue=False, + # --- inner event handlers --------------------------------------------------- + + def _ui_toggle_catalog(enabled: bool): + self._use_tag_catalog = bool(enabled) + if not self._use_tag_catalog: + self._set_catalog_source("bundled") + ok, message = self._load_tag_catalog() + if not ok: + self._catalog = NoopCatalog() + self._tag_catalog_status_text = message + self._save_tag_catalog_preferences() + return ( + _gr_component_update( + gr.Radio, + visible=self._use_tag_catalog, + value=("Custom file" if self._catalog_source == "custom" else "Bundled"), + ), + _gr_component_update( + gr.Group, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + ), + _gr_component_update( + gr.Textbox, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + value=self._custom_catalog_path, + ), + _gr_component_update(gr.Button, visible=self._use_tag_catalog), + _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), ) - gelbooru_save_button.click( - fn=self._ui_save_gelbooru_credentials, - inputs=[gelbooru_api_key, gelbooru_user_id], - outputs=[ - gelbooru_saved_message, - gelbooru_credentials_group, - gelbooru_clear_button, - gelbooru_api_key, - gelbooru_user_id, - ], - queue=False, + + def _ui_set_catalog_source(source_label: str): + source = "custom" if (source_label or "") == "Custom file" else "bundled" + self._set_catalog_source(source) + if self._use_tag_catalog: + ok, message = self._load_tag_catalog() + if not ok: + self._catalog = NoopCatalog() + self._tag_catalog_status_text = message + else: + self._tag_catalog_status_text = self._format_catalog_status() + else: + self._tag_catalog_status_text = self._format_catalog_status() + self._save_tag_catalog_preferences() + self._update_tag_diag() + return ( + _gr_component_update( + gr.Group, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + ), + _gr_component_update( + gr.Textbox, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + value=self._custom_catalog_path, + ), + _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), ) - gelbooru_clear_button.click( - fn=self._ui_clear_gelbooru_credentials, - inputs=[], - outputs=[ - gelbooru_saved_message, - gelbooru_credentials_group, - gelbooru_clear_button, - gelbooru_api_key, - gelbooru_user_id, - ], - queue=False, + + def _ui_set_catalog_path(path_value: str): + self._custom_catalog_path = (path_value or "").strip() + self._tag_catalog_path = self._custom_catalog_path + if self._use_tag_catalog and self._catalog_source == "custom": + if self._custom_catalog_path: + ok, message = self._load_tag_catalog() + if not ok: + self._catalog = NoopCatalog() + self._tag_catalog_status_text = message + else: + self._tag_catalog_status_text = self._format_catalog_status() + else: + self._catalog = NoopCatalog() + self._tag_catalog_status_text = "Catalog mode: ON - No path set" + else: + self._tag_catalog_status_text = self._format_catalog_status() + self._save_tag_catalog_preferences() + self._update_tag_diag() + return ( + _gr_component_update( + gr.Textbox, + value=self._custom_catalog_path, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + ), + _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), ) - gr.Markdown("""\n---\n""") - with gr.Group(): - with gr.Accordion("Img2Img / ControlNet", open=False): - use_img2img = gr.Checkbox(label="Use Image for Img2Img", value=False) - use_ip = gr.Checkbox(label="Use Image for ControlNet (Unit 0)", value=False) - denoising = gr.Slider( - value=0.75, - label="Img2Img Denoising / CN Weight", - minimum=0.0, - maximum=1.0, - step=0.05, - ) - use_last_img = gr.Checkbox(label="Use same image for batch", value=False) - crop_center = gr.Checkbox(label="Crop image to fit target", value=False) - enable_adetailer_support = gr.Checkbox( - label="Enable RanbooruX ADetailer support", - value=False, - info="Run RanbooruX's manual ADetailer integration after img2img when enabled.", - ) - reuse_cached_posts = gr.Checkbox( - label="Reuse cached booru posts", - value=False, - info="Leave disabled to fetch fresh images every generation. Enable when you want RanbooruX to reuse the previously cached posts.", - ) - with gr.Group(): - with gr.Accordion("File Tags", open=False): - use_search_txt = gr.Checkbox(label="Add line from Search File", value=False) - choose_search_txt = gr.Dropdown( - self.get_files(USER_SEARCH_DIR), - label="Choose Search File", - value="", - info=f"in '{USER_SEARCH_DIR}'", - ) - search_refresh_btn = gr.Button("Refresh") - use_remove_txt = gr.Checkbox(label="Add tags from Remove File", value=False) - choose_remove_txt = gr.Dropdown( - self.get_files(USER_REMOVE_DIR), - label="Choose Remove File", - value="", - info=f"in '{USER_REMOVE_DIR}'", - ) - remove_refresh_btn = gr.Button("Refresh") - with gr.Group(): - with gr.Accordion("Extra Prompt Modes", open=False): - with gr.Box(): - mix_prompt = gr.Checkbox(label="Mix tags from multiple posts", value=False) - mix_amount = gr.Slider( - value=2, label="Posts to mix", minimum=2, maximum=10, step=1 - ) - with gr.Box(): - chaos_mode = gr.Radio( - ["None", "Shuffle All", "Shuffle Negative"], - label="Tag Shuffling (Chaos)", - value="None", - ) - chaos_amount = gr.Slider( - value=0.5, label="Chaos Amount %", minimum=0.1, maximum=1.0, step=0.05 - ) - with gr.Box(): - use_same_seed = gr.Checkbox(label="Use same seed for batch", value=False) - use_cache = gr.Checkbox(label="Cache Booru API requests", value=True) - log_prompt_sources = gr.Checkbox( - label="Log image sources/prompts to txt", - value=False, - info="When enabled, RanbooruX appends a log entry mapping seeds and prompts to the source posts.", - ) + def _ui_reload_catalog(): + if self._use_tag_catalog: + ok, message = self._load_tag_catalog() + if not ok: + self._catalog = NoopCatalog() + self._tag_catalog_status_text = message + else: + self._tag_catalog_status_text = self._format_catalog_status() + else: + self._tag_catalog_status_text = self._format_catalog_status() + self._save_tag_catalog_preferences() + self._update_tag_diag() + return _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text) + + def _ui_catalog_upload(uploaded): + guessed_path = self._catalog_path_from_upload(uploaded) + if guessed_path: + self._custom_catalog_path = guessed_path + self._tag_catalog_path = guessed_path + self._save_tag_catalog_preferences() + msg = f"Selected custom catalog file: {os.path.basename(guessed_path)}" + else: + msg = self._tag_catalog_status_text + return ( + _gr_component_update( + gr.Textbox, + value=self._custom_catalog_path, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + ), + _gr_component_update(gr.Markdown, value=msg), + ) + + def _ui_validate_catalog(path_value, uploaded): + candidate = (path_value or "").strip() or self._catalog_path_from_upload(uploaded) + ok, message = self._validate_csv_format(candidate) + status = f"Validation passed: {message}" if ok else f"Validation failed: {message}" + return _gr_component_update(gr.Markdown, value=status) + + def _ui_import_custom_catalog(uploaded, path_value): + ok, message = self._import_custom_catalog(uploaded, path_hint=path_value) + if not ok: + return ( + _gr_component_update( + gr.Radio, + value=("Custom file" if self._catalog_source == "custom" else "Bundled"), + ), + _gr_component_update( + gr.Group, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + ), + _gr_component_update( + gr.Textbox, + value=self._custom_catalog_path, + visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), + ), + _gr_component_update(gr.Markdown, value=message), + ) + self._tag_catalog_status_text = self._format_catalog_status() + self._update_tag_diag() + return ( + _gr_component_update(gr.Radio, value="Custom file"), + _gr_component_update(gr.Group, visible=True), + _gr_component_update(gr.Textbox, value=self._custom_catalog_path, visible=True), + _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), + ) + + # --- event wiring ----------------------------------------------------------- + + use_tag_catalog.change( + fn=_ui_toggle_catalog, + inputs=[use_tag_catalog], + outputs=[ + catalog_source, + custom_catalog_group, + catalog_path, + reload_catalog, + catalog_status, + ], + queue=False, + ) + catalog_source.change( + fn=_ui_set_catalog_source, + inputs=[catalog_source], + outputs=[custom_catalog_group, catalog_path, catalog_status], + queue=False, + ) + catalog_path.change( + fn=_ui_set_catalog_path, + inputs=[catalog_path], + outputs=[catalog_path, catalog_status], + queue=False, + ) + catalog_upload.upload( + fn=_ui_catalog_upload, + inputs=[catalog_upload], + outputs=[catalog_path, catalog_status], + queue=False, + ) + catalog_validate_btn.click( + fn=_ui_validate_catalog, + inputs=[catalog_path, catalog_upload], + outputs=[catalog_status], + queue=False, + ) + catalog_import_btn.click( + fn=_ui_import_custom_catalog, + inputs=[catalog_upload, catalog_path], + outputs=[catalog_source, custom_catalog_group, catalog_path, catalog_status], + queue=False, + ) + reload_catalog.click( + fn=_ui_reload_catalog, + inputs=None, + outputs=[catalog_status], + queue=False, + ) + + # --- Platform Diagnostics --------------------------------------------------- + + diagnostics_visible_state = gr.State(False) + diagnostics_toggle_btn = gr.Button("Show Platform Diagnostics") + diagnostics_md = gr.Markdown("", visible=False) + diagnostics_toggle_btn.click( + fn=self._toggle_platform_diagnostics, + inputs=[diagnostics_visible_state], + outputs=[diagnostics_visible_state, diagnostics_md, diagnostics_toggle_btn], + queue=False, + ) + + return use_tag_catalog, catalog_path + + def _build_lora_ui_section(self): + """LoRAnado controls, auto-detect, detected LoRAs, and blacklist. + + Performs the initial LoRA scan, creates all LoRAnado widgets inside + ``InputAccordion``, and wires up the change/click events. Returns the + components that must appear in the script-args component list. + """ initial_lora_scan = self._scan_loranado_candidates("") initial_lora_choices = ( initial_lora_scan.get("detected_names") or initial_lora_scan.get("all_names") or [] @@ -3184,10 +2035,10 @@ def ui(self, is_img2img): if initial_lora_scan.get("all_names"): if initial_lora_scan.get("detected_names"): initial_lora_status = ( - f"Detected {len(initial_lora_scan['detected_names'])} PonyXL/Anima-compatible LoRAs." + f"Detected {len(initial_lora_scan['detected_names'])} PonyXL-compatible LoRAs." ) else: - initial_lora_status = f"No PonyXL/Anima markers detected; using all {len(initial_lora_scan['all_names'])} LoRAs." + initial_lora_status = f"No PonyXL markers detected; using all {len(initial_lora_scan['all_names'])} LoRAs." with InputAccordion( False, label="LoRAnado", elem_id=self.elem_id("lo_enable") @@ -3212,9 +2063,9 @@ def ui(self, is_img2img): ) with gr.Box(): lora_auto_detect_pony = gr.Checkbox( - label="Auto-detect PonyXL/Anima-compatible LoRAs", + label="Auto-detect PonyXL-compatible LoRAs", value=True, - info="Scans LoRA filenames and safetensors metadata for PonyXL and Anima markers.", + info="Scans LoRA filenames and safetensors metadata for PonyXL markers.", ) with gr.Row(): lora_scan_btn = gr.Button("Scan LoRAs") @@ -3234,335 +2085,549 @@ def ui(self, is_img2img): info="Blacklisted LoRAs are excluded from random selection.", ) lora_detect_status = gr.Markdown(initial_lora_status) - search_refresh_btn.click(fn=self.refresh_ser, inputs=[], outputs=[choose_search_txt]) - remove_refresh_btn.click(fn=self.refresh_rem, inputs=[], outputs=[choose_remove_txt]) - personal_add_btn.click( - fn=self._ui_add_personal_tags, - inputs=[personal_remove_input, personal_remove_dropdown], - outputs=[personal_remove_dropdown, personal_remove_input], + + # --- LoRA event wiring ---------------------------------------------------- + + lora_folder.change( + fn=self._ui_refresh_loranado_controls, + inputs=[lora_folder, lora_auto_detect_pony, lora_detected_loras, lora_blacklist], + outputs=[lora_detected_loras, lora_blacklist, lora_detect_status], queue=False, ) - personal_remove_btn.click( - fn=self._ui_remove_personal_tags, - inputs=[personal_remove_dropdown], - outputs=[personal_remove_dropdown], + lora_auto_detect_pony.change( + fn=self._ui_refresh_loranado_controls, + inputs=[lora_folder, lora_auto_detect_pony, lora_detected_loras, lora_blacklist], + outputs=[lora_detected_loras, lora_blacklist, lora_detect_status], queue=False, ) - personal_dedupe_btn.click( - fn=self._ui_dedupe_personal_list, - inputs=[], - outputs=[personal_remove_dropdown], + lora_scan_btn.click( + fn=self._ui_refresh_loranado_controls, + inputs=[lora_folder, lora_auto_detect_pony, lora_detected_loras, lora_blacklist], + outputs=[lora_detected_loras, lora_blacklist, lora_detect_status], queue=False, ) - personal_import_file.upload( - fn=self._ui_import_personal_list, - inputs=[personal_import_file], - outputs=[personal_remove_dropdown, personal_import_file], + lora_select_all_btn.click( + fn=self._ui_select_all_loranado, + inputs=[lora_folder, lora_auto_detect_pony, lora_blacklist], + outputs=[lora_detected_loras, lora_detect_status], queue=False, ) - personal_export_btn.click( - fn=self._ui_export_personal_list, inputs=[], outputs=None, queue=False - ) - favorites_add_btn.click( - fn=self._ui_add_favorite_tags, - inputs=[favorites_input, favorites_dropdown], - outputs=[favorites_dropdown, favorites_input], - queue=False, + return ( + lora_enabled, + lora_folder, + lora_amount, + lora_min, + lora_max, + lora_custom_weights, + lora_lock_prev, + lora_auto_detect_pony, + lora_detected_loras, + lora_blacklist, ) - favorites_remove_btn.click( - fn=self._ui_remove_favorite_tags, - inputs=[favorites_dropdown], - outputs=[favorites_dropdown], + + def _build_filter_ui_section(self): + """Removal toggle checkboxes, presets, and Quick Strip. + + Creates the Quick Presets buttons and all removal-filter checkboxes + (Text & Metadata, Characters & Series, Clothing, Furry & Headwear, + Girl Suffix, Colors & Traits, Subject Constraints). Wires up the + preset click events. Must be called inside ``gr.Accordion("Removal Filters")`` + after the catalog section. Returns the 11 filter components that + appear in the script-args component list. + """ + gr.Markdown("**Quick Presets**: apply common filter combinations with one click.") + + with gr.Row(): + preset_strip_series = gr.Button("Strip Series/Character") + preset_remove_text = gr.Button("Remove Text-like Tags") + preset_preserve_colors = gr.Button("Preserve Base Colors") + preset_quick_strip = gr.Button("Quick Strip") + with gr.Group(): + gr.Markdown("**Text & Metadata**") + remove_bad_tags = gr.Checkbox( + label="Remove common 'bad' tags", + value=True, + info="Cull frequent watermark, commentary, and UI text tags from prompts.", + ) + remove_text_tags = gr.Checkbox( + label="Remove tag/text/commentary metadata", + value=True, + info="Strip speech bubbles, watermark text, and similar metadata from fetched prompts.", + ) + with gr.Group(): + gr.Markdown("**Characters & Series**") + remove_artist_tags = gr.Checkbox( + label="Remove artist tags", + value=False, + info="Drop artist credits drawn from the source post.", + ) + remove_character_tags = gr.Checkbox( + label="Remove character tags", + value=False, + info="Filter character/franchise tags sourced from metadata.", + ) + remove_series_tags = gr.Checkbox( + label="Remove series / franchise tags", + value=False, + info="Ignore franchise/game/anime tags to keep prompts generic.", + ) + with gr.Group(): + gr.Markdown("**Clothing & Accessories**") + remove_clothing_tags = gr.Checkbox( + label="Remove clothing tags", + value=False, + info="Omit apparel/accessory tags introduced by the booru.", + ) + with gr.Group(): + gr.Markdown("**Furry & Headwear**") + remove_furry_tags = gr.Checkbox( + label="Filter furry/pokemon tags", + value=False, + info="Remove furry, pokemon, and animal trait tags.", + ) + remove_headwear_tags = gr.Checkbox( + label="Filter headwear / halo tags", + value=False, + info="Strip hats, halos, and similar head accessories.", + ) + with gr.Group(): + gr.Markdown("**Girl Suffix**") + remove_girl_suffix_tags = gr.Checkbox( + label="Filter _girl suffix tags", + value=False, + info="Remove demon_girl, cat_girl, angel_girl and similar *_girl tags (keeps 1girl, 2girls, etc.).", + ) + with gr.Group(): + + gr.Markdown("**Colors & Traits**") + preserve_hair_eye_colors = gr.Checkbox( + label="Preserve base hair & eye colors", + value=False, + info="Keep your prompt's hair/eye colors while removing conflicting imports.", + ) + with gr.Group(): + gr.Markdown("**Subject Constraints**") + restrict_subject_tags = gr.Checkbox( + label="Keep only subject counts", + value=False, + info="Maintain your subject count (e.g., solo/1girl) by removing mismatched tags.", + ) + + # --- preset wiring --------------------------------------------------------- + + preset_strip_series.click( + fn=lambda: ( + _gr_component_update(gr.Checkbox, value=True), + _gr_component_update(gr.Checkbox, value=True), + _gr_component_update(gr.Checkbox, value=True), + ), + inputs=[], + outputs=[remove_series_tags, remove_character_tags, remove_artist_tags], queue=False, ) - favorites_dedupe_btn.click( - fn=self._ui_dedupe_favorite_list, inputs=[], outputs=[favorites_dropdown], queue=False + preset_remove_text.click( + fn=lambda: ( + _gr_component_update(gr.Checkbox, value=True), + _gr_component_update(gr.Checkbox, value=True), + ), + inputs=[], + outputs=[remove_text_tags, remove_bad_tags], + queue=False, ) - favorites_import_file.upload( - fn=self._ui_import_favorite_list, - inputs=[favorites_import_file], - outputs=[favorites_dropdown, favorites_import_file], + preset_preserve_colors.click( + fn=lambda: _gr_component_update(gr.Checkbox, value=True), + inputs=[], + outputs=[preserve_hair_eye_colors], queue=False, ) - favorites_export_btn.click( - fn=self._ui_export_favorite_list, inputs=[], outputs=None, queue=False + preset_quick_strip.click( + fn=lambda: tuple(_gr_component_update(gr.Checkbox, value=True) for _ in range(11)), + inputs=[], + outputs=[ + remove_bad_tags, + remove_text_tags, + remove_artist_tags, + remove_character_tags, + remove_series_tags, + remove_clothing_tags, + remove_furry_tags, + remove_headwear_tags, + remove_girl_suffix_tags, + preserve_hair_eye_colors, + restrict_subject_tags, + ], + queue=False, ) - def _ui_toggle_catalog(enabled: bool): - self._use_tag_catalog = bool(enabled) - if not self._use_tag_catalog: - self._set_catalog_source("bundled") - ok, message = self._load_tag_catalog() - if not ok: - self._catalog = NoopCatalog() - self._tag_catalog_status_text = message - self._save_tag_catalog_preferences() - return ( - _gr_component_update( - gr.Radio, - visible=self._use_tag_catalog, - value=("Custom file" if self._catalog_source == "custom" else "Bundled"), - ), - _gr_component_update( - gr.Group, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - ), - _gr_component_update( - gr.Textbox, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - value=self._custom_catalog_path, - ), - _gr_component_update(gr.Button, visible=self._use_tag_catalog), - _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), - ) + return ( + remove_bad_tags, + remove_text_tags, + remove_artist_tags, + remove_character_tags, + remove_series_tags, + remove_clothing_tags, + remove_furry_tags, + remove_headwear_tags, + remove_girl_suffix_tags, + preserve_hair_eye_colors, + restrict_subject_tags, + ) - def _ui_set_catalog_source(source_label: str): - source = "custom" if (source_label or "") == "Custom file" else "bundled" - self._set_catalog_source(source) - if self._use_tag_catalog: - ok, message = self._load_tag_catalog() - if not ok: - self._catalog = NoopCatalog() - self._tag_catalog_status_text = message - else: - self._tag_catalog_status_text = self._format_catalog_status() - else: - self._tag_catalog_status_text = self._format_catalog_status() - self._save_tag_catalog_preferences() - self._update_tag_diag() - return ( - _gr_component_update( - gr.Group, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - ), - _gr_component_update( - gr.Textbox, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - value=self._custom_catalog_path, - ), - _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), + def _build_personal_lists_ui_section(self): + """Search/remove file management with refresh buttons (File Tags accordion). + + Creates the File Tags accordion containing search-file and remove-file + dropdowns with Refresh buttons. Wires the refresh click events. Returns + the six components needed in the script-args list. + """ + with gr.Accordion("File Tags", open=False): + use_search_txt = gr.Checkbox(label="Add line from Search File", value=False) + choose_search_txt = gr.Dropdown( + self.get_files(USER_SEARCH_DIR), + label="Choose Search File", + value="", + info=f"in '{USER_SEARCH_DIR}'", ) - - def _ui_set_catalog_path(path_value: str): - self._custom_catalog_path = (path_value or "").strip() - self._tag_catalog_path = self._custom_catalog_path - if self._use_tag_catalog and self._catalog_source == "custom": - if self._custom_catalog_path: - ok, message = self._load_tag_catalog() - if not ok: - self._catalog = NoopCatalog() - self._tag_catalog_status_text = message - else: - self._tag_catalog_status_text = self._format_catalog_status() - else: - self._catalog = NoopCatalog() - self._tag_catalog_status_text = "Catalog mode: ON - No path set" - else: - self._tag_catalog_status_text = self._format_catalog_status() - self._save_tag_catalog_preferences() - self._update_tag_diag() - return ( - _gr_component_update( - gr.Textbox, - value=self._custom_catalog_path, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - ), - _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), + search_refresh_btn = gr.Button("Refresh") + use_remove_txt = gr.Checkbox(label="Add tags from Remove File", value=False) + choose_remove_txt = gr.Dropdown( + self.get_files(USER_REMOVE_DIR), + label="Choose Remove File", + value="", + info=f"in '{USER_REMOVE_DIR}'", ) + remove_refresh_btn = gr.Button("Refresh") - def _ui_reload_catalog(): - if self._use_tag_catalog: - ok, message = self._load_tag_catalog() - if not ok: - self._catalog = NoopCatalog() - self._tag_catalog_status_text = message - else: - self._tag_catalog_status_text = self._format_catalog_status() - else: - self._tag_catalog_status_text = self._format_catalog_status() - self._save_tag_catalog_preferences() - self._update_tag_diag() - return _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text) - - def _ui_catalog_upload(uploaded): - guessed_path = self._catalog_path_from_upload(uploaded) - if guessed_path: - self._custom_catalog_path = guessed_path - self._tag_catalog_path = guessed_path - self._save_tag_catalog_preferences() - msg = f"Selected custom catalog file: {os.path.basename(guessed_path)}" - else: - msg = self._tag_catalog_status_text - return ( - _gr_component_update( - gr.Textbox, - value=self._custom_catalog_path, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - ), - _gr_component_update(gr.Markdown, value=msg), - ) + search_refresh_btn.click(fn=self.refresh_ser, inputs=[], outputs=[choose_search_txt]) + remove_refresh_btn.click(fn=self.refresh_rem, inputs=[], outputs=[choose_remove_txt]) - def _ui_validate_catalog(path_value, uploaded): - candidate = (path_value or "").strip() or self._catalog_path_from_upload(uploaded) - ok, message = self._validate_csv_format(candidate) - status = f"Validation passed: {message}" if ok else f"Validation failed: {message}" - return _gr_component_update(gr.Markdown, value=status) + return ( + use_search_txt, + use_remove_txt, + choose_search_txt, + choose_remove_txt, + search_refresh_btn, + remove_refresh_btn, + ) - def _ui_import_custom_catalog(uploaded, path_value): - ok, message = self._import_custom_catalog(uploaded, path_hint=path_value) - if not ok: - return ( - _gr_component_update( - gr.Radio, - value=("Custom file" if self._catalog_source == "custom" else "Bundled"), - ), - _gr_component_update( - gr.Group, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - ), - _gr_component_update( - gr.Textbox, - value=self._custom_catalog_path, - visible=bool(self._use_tag_catalog and self._catalog_source == "custom"), - ), - _gr_component_update(gr.Markdown, value=message), + def ui(self, is_img2img): + with InputAccordion(False, label="RanbooruX", elem_id=self.elem_id("ra_enable")) as enabled: + booru_list = [ + "danbooru", + "gelbooru", + "gelbooru-compatible", + "xbooru", + "rule34", + "safebooru", + "konachan", + "yande.re", + "aibooru", + "e621", + ] + booru = gr.Dropdown(booru_list, label="Booru", value="danbooru") + with gr.Group(visible=False) as gelbooru_credentials_group: + gelbooru_api_key = gr.Textbox( + label="Gelbooru API Key", + type="password", + placeholder="Enter your Gelbooru API key", ) - self._tag_catalog_status_text = self._format_catalog_status() - self._update_tag_diag() - return ( - _gr_component_update(gr.Radio, value="Custom file"), - _gr_component_update(gr.Group, visible=True), - _gr_component_update(gr.Textbox, value=self._custom_catalog_path, visible=True), - _gr_component_update(gr.Markdown, value=self._tag_catalog_status_text), + gelbooru_user_id = gr.Textbox( + label="Gelbooru User ID", placeholder="Enter your Gelbooru user ID" + ) + gelbooru_save_button = gr.Button("Save Credentials to Disk", variant="primary") + gelbooru_saved_message = gr.Markdown("", visible=False) + gelbooru_clear_button = gr.Button("Clear Saved Credentials", visible=False) + with gr.Group(visible=False) as gelbooru_compat_group: + gelbooru_compat_base_url = gr.Textbox( + label="Gelbooru-compatible Base URL", + placeholder="https://realbooru.com", + value=self._gelbooru_compat_base_url, + ) + max_pages = gr.Slider( + label="Max Pages (tag search)", minimum=1, maximum=100, value=10, step=1 + ) + gr.Markdown("""## Post""") + post_id = gr.Textbox(lines=1, label="Post ID (Overrides tags/pages)") + gr.Markdown("""## Tags""") + tags = gr.Textbox(lines=1, label="Tags to Search (Pre)") + remove_tags = gr.Textbox(lines=1, label="Tags to Remove (Post)") + mature_rating = gr.Radio( + list(RATINGS.get("gelbooru", RATING_TYPES["none"])), + label="Mature Rating", + value="All", + ) + with gr.Accordion("Removal Filters", open=False): + with gr.Group(): + use_tag_catalog, catalog_path = self._build_catalog_ui_section() + + ( + remove_bad_tags, + remove_text_tags, + remove_artist_tags, + remove_character_tags, + remove_series_tags, + remove_clothing_tags, + remove_furry_tags, + remove_headwear_tags, + remove_girl_suffix_tags, + preserve_hair_eye_colors, + restrict_subject_tags, + ) = self._build_filter_ui_section() + personal_choices = self._read_list_file(PERSONAL_REMOVE_FILE) + favorite_choices = self._read_list_file(FAVORITES_FILE) + with gr.Accordion("Personal Lists", open=False): + with gr.Row(): + with gr.Column(): + gr.Markdown("**Personal Removal List**") + personal_remove_dropdown = gr.Dropdown( + choices=personal_choices, + value=personal_choices, + multiselect=True, + label="Removal Tags", + allow_custom_value=False, + ) + personal_remove_input = gr.Textbox( + label="Add tags", placeholder="comma or newline separated" + ) + with gr.Row(): + personal_add_btn = gr.Button("Add", variant="primary") + personal_remove_btn = gr.Button("Remove Selected") + personal_dedupe_btn = gr.Button("De-duplicate") + with gr.Row(): + personal_import_file = gr.File( + label="Import CSV/TXT", file_types=[".txt", ".csv"], visible=True + ) + personal_export_btn = gr.DownloadButton("Export") + with gr.Column(): + gr.Markdown("**Favorites List**") + favorites_dropdown = gr.Dropdown( + choices=favorite_choices, + value=favorite_choices, + multiselect=True, + label="Favorite Tags", + allow_custom_value=False, + ) + favorites_input = gr.Textbox( + label="Add favorites", placeholder="comma or newline separated" + ) + with gr.Row(): + favorites_add_btn = gr.Button("Add", variant="primary") + favorites_remove_btn = gr.Button("Remove Selected") + favorites_dedupe_btn = gr.Button("De-duplicate") + with gr.Row(): + favorites_import_file = gr.File( + label="Import CSV/TXT", file_types=[".txt", ".csv"], visible=True + ) + favorites_export_btn = gr.DownloadButton("Export") + shuffle_tags = gr.Checkbox(label="Shuffle tags", value=True) + change_dash = gr.Checkbox(label='Convert "_" to spaces', value=False) + anima_auto_detect = gr.Checkbox( + label="Auto-detect Anima model", + value=True, + info="Automatically enable space-separated tags when an Anima model is loaded", + ) + anima_tune_img2img = gr.Checkbox( + label="Auto-tune Img2Img parameters for Anima", + value=True, + info="Automatically optimize steps, CFG scale, and denoising for Anima flow-matching", + ) + same_prompt = gr.Checkbox(label="Use same prompt for batch", value=False) + fringe_benefits = gr.Checkbox( + label="Gelbooru: Fringe Benefits", value=True, visible=False + ) + limit_tags = gr.Slider( + value=1.0, label="Limit tags by %", minimum=0.05, maximum=1.0, step=0.05 + ) + max_tags = gr.Slider( + value=0, label="Max tags (0=disabled)", minimum=0, maximum=300, step=1 + ) + change_background = gr.Radio( + ["Don't Change", "Add Detail", "Force Simple", "Force Transparent/White"], + label="Change Background", + value="Don't Change", + ) + change_color = gr.Radio( + ["Don't Change", "Force Color", "Force Monochrome"], + label="Change Color", + value="Don't Change", + ) + sorting_order = gr.Radio( + ["Random", "Score Descending", "Score Ascending"], + label="Sort Order (tag search)", + value="Random", + ) + booru.change(get_available_ratings, booru, mature_rating) + booru.change(show_fringe_benefits, booru, fringe_benefits) + booru.change( + self._update_gelbooru_ui_visibility, + inputs=[booru], + outputs=[ + gelbooru_credentials_group, + gelbooru_saved_message, + gelbooru_clear_button, + gelbooru_api_key, + gelbooru_user_id, + ], + queue=False, + ) + booru.change( + self._update_gelbooru_compat_visibility, + inputs=[booru], + outputs=[gelbooru_compat_group, gelbooru_compat_base_url], + queue=False, + ) + gelbooru_compat_base_url.change( + fn=self._ui_set_gelbooru_compat_base_url, + inputs=[gelbooru_compat_base_url], + outputs=[gelbooru_compat_base_url], + queue=False, + ) + gelbooru_save_button.click( + fn=self._ui_save_gelbooru_credentials, + inputs=[gelbooru_api_key, gelbooru_user_id], + outputs=[ + gelbooru_saved_message, + gelbooru_credentials_group, + gelbooru_clear_button, + gelbooru_api_key, + gelbooru_user_id, + ], + queue=False, + ) + gelbooru_clear_button.click( + fn=self._ui_clear_gelbooru_credentials, + inputs=[], + outputs=[ + gelbooru_saved_message, + gelbooru_credentials_group, + gelbooru_clear_button, + gelbooru_api_key, + gelbooru_user_id, + ], + queue=False, ) - use_tag_catalog.change( - fn=_ui_toggle_catalog, - inputs=[use_tag_catalog], - outputs=[ - catalog_source, - custom_catalog_group, - catalog_path, - reload_catalog, - catalog_status, - ], - queue=False, - ) - catalog_source.change( - fn=_ui_set_catalog_source, - inputs=[catalog_source], - outputs=[custom_catalog_group, catalog_path, catalog_status], - queue=False, - ) - catalog_path.change( - fn=_ui_set_catalog_path, - inputs=[catalog_path], - outputs=[catalog_path, catalog_status], - queue=False, - ) - catalog_upload.upload( - fn=_ui_catalog_upload, - inputs=[catalog_upload], - outputs=[catalog_path, catalog_status], - queue=False, - ) - catalog_validate_btn.click( - fn=_ui_validate_catalog, - inputs=[catalog_path, catalog_upload], - outputs=[catalog_status], - queue=False, - ) - catalog_import_btn.click( - fn=_ui_import_custom_catalog, - inputs=[catalog_upload, catalog_path], - outputs=[catalog_source, custom_catalog_group, catalog_path, catalog_status], - queue=False, - ) - reload_catalog.click( - fn=_ui_reload_catalog, - inputs=None, - outputs=[catalog_status], + gr.Markdown("""\n---\n""") + with gr.Group(): + with gr.Accordion("Img2Img / ControlNet", open=False): + use_img2img = gr.Checkbox(label="Use Image for Img2Img", value=False) + use_ip = gr.Checkbox(label="Use Image for ControlNet (Unit 0)", value=False) + denoising = gr.Slider( + value=0.75, + label="Img2Img Denoising / CN Weight", + minimum=0.0, + maximum=1.0, + step=0.05, + ) + use_last_img = gr.Checkbox(label="Use same image for batch", value=False) + crop_center = gr.Checkbox(label="Crop image to fit target", value=False) + enable_adetailer_support = gr.Checkbox( + label="Enable RanbooruX ADetailer support", + value=False, + info="Run RanbooruX's manual ADetailer integration after img2img when enabled.", + ) + reuse_cached_posts = gr.Checkbox( + label="Reuse cached booru posts", + value=False, + info="Leave disabled to fetch fresh images every generation. Enable when you want RanbooruX to reuse the previously cached posts.", + ) + with gr.Group(): + ( + use_search_txt, + use_remove_txt, + choose_search_txt, + choose_remove_txt, + search_refresh_btn, + remove_refresh_btn, + ) = self._build_personal_lists_ui_section() + with gr.Group(): + with gr.Accordion("Extra Prompt Modes", open=False): + with gr.Box(): + mix_prompt = gr.Checkbox(label="Mix tags from multiple posts", value=False) + mix_amount = gr.Slider( + value=2, label="Posts to mix", minimum=2, maximum=10, step=1 + ) + with gr.Box(): + chaos_mode = gr.Radio( + ["None", "Shuffle All", "Shuffle Negative"], + label="Tag Shuffling (Chaos)", + value="None", + ) + chaos_amount = gr.Slider( + value=0.5, label="Chaos Amount %", minimum=0.1, maximum=1.0, step=0.05 + ) + with gr.Box(): + use_same_seed = gr.Checkbox(label="Use same seed for batch", value=False) + use_cache = gr.Checkbox(label="Cache Booru API requests", value=True) + log_prompt_sources = gr.Checkbox( + label="Log image sources/prompts to txt", + value=False, + info="When enabled, RanbooruX appends a log entry mapping seeds and prompts to the source posts.", + ) + ( + lora_enabled, + lora_folder, + lora_amount, + lora_min, + lora_max, + lora_custom_weights, + lora_lock_prev, + lora_auto_detect_pony, + lora_detected_loras, + lora_blacklist, + ) = self._build_lora_ui_section() + personal_add_btn.click( + fn=self._ui_add_personal_tags, + inputs=[personal_remove_input, personal_remove_dropdown], + outputs=[personal_remove_dropdown, personal_remove_input], queue=False, ) - - preset_strip_series.click( - fn=lambda: ( - _gr_component_update(gr.Checkbox, value=True), - _gr_component_update(gr.Checkbox, value=True), - _gr_component_update(gr.Checkbox, value=True), - ), - inputs=[], - outputs=[remove_series_tags, remove_character_tags, remove_artist_tags], + personal_remove_btn.click( + fn=self._ui_remove_personal_tags, + inputs=[personal_remove_dropdown], + outputs=[personal_remove_dropdown], queue=False, ) - preset_remove_text.click( - fn=lambda: ( - _gr_component_update(gr.Checkbox, value=True), - _gr_component_update(gr.Checkbox, value=True), - ), + personal_dedupe_btn.click( + fn=self._ui_dedupe_personal_list, inputs=[], - outputs=[remove_text_tags, remove_bad_tags], + outputs=[personal_remove_dropdown], queue=False, ) - preset_preserve_colors.click( - fn=lambda: _gr_component_update(gr.Checkbox, value=True), - inputs=[], - outputs=[preserve_hair_eye_colors], + personal_import_file.upload( + fn=self._ui_import_personal_list, + inputs=[personal_import_file], + outputs=[personal_remove_dropdown, personal_import_file], queue=False, ) - preset_quick_strip.click( - fn=lambda: tuple(_gr_component_update(gr.Checkbox, value=True) for _ in range(11)), - inputs=[], - outputs=[ - remove_bad_tags, - remove_text_tags, - remove_artist_tags, - remove_character_tags, - remove_series_tags, - remove_clothing_tags, - remove_furry_tags, - remove_headwear_tags, - remove_girl_suffix_tags, - preserve_hair_eye_colors, - restrict_subject_tags, - ], - queue=False, + personal_export_btn.click( + fn=self._ui_export_personal_list, inputs=[], outputs=None, queue=False ) - lora_folder.change( - fn=self._ui_refresh_loranado_controls, - inputs=[lora_folder, lora_auto_detect_pony, lora_detected_loras, lora_blacklist], - outputs=[lora_detected_loras, lora_blacklist, lora_detect_status], + favorites_add_btn.click( + fn=self._ui_add_favorite_tags, + inputs=[favorites_input, favorites_dropdown], + outputs=[favorites_dropdown, favorites_input], queue=False, ) - lora_auto_detect_pony.change( - fn=self._ui_refresh_loranado_controls, - inputs=[lora_folder, lora_auto_detect_pony, lora_detected_loras, lora_blacklist], - outputs=[lora_detected_loras, lora_blacklist, lora_detect_status], + favorites_remove_btn.click( + fn=self._ui_remove_favorite_tags, + inputs=[favorites_dropdown], + outputs=[favorites_dropdown], queue=False, ) - lora_scan_btn.click( - fn=self._ui_refresh_loranado_controls, - inputs=[lora_folder, lora_auto_detect_pony, lora_detected_loras, lora_blacklist], - outputs=[lora_detected_loras, lora_blacklist, lora_detect_status], - queue=False, + favorites_dedupe_btn.click( + fn=self._ui_dedupe_favorite_list, inputs=[], outputs=[favorites_dropdown], queue=False ) - lora_select_all_btn.click( - fn=self._ui_select_all_loranado, - inputs=[lora_folder, lora_auto_detect_pony, lora_blacklist], - outputs=[lora_detected_loras, lora_detect_status], + favorites_import_file.upload( + fn=self._ui_import_favorite_list, + inputs=[favorites_import_file], + outputs=[favorites_dropdown, favorites_import_file], queue=False, ) - - diagnostics_visible_state = gr.State(False) - diagnostics_toggle_btn = gr.Button("Show Platform Diagnostics") - diagnostics_md = gr.Markdown("", visible=False) - diagnostics_toggle_btn.click( - fn=self._toggle_platform_diagnostics, - inputs=[diagnostics_visible_state], - outputs=[diagnostics_visible_state, diagnostics_md, diagnostics_toggle_btn], - queue=False, + favorites_export_btn.click( + fn=self._ui_export_favorite_list, inputs=[], outputs=None, queue=False ) components = [ @@ -3628,6 +2693,8 @@ def _ui_import_custom_catalog(uploaded, path_value): lora_auto_detect_pony, lora_detected_loras, lora_blacklist, + anima_auto_detect, + anima_tune_img2img, ] return rb_run_options.RunComponents.from_sequence(components).script_args() @@ -3670,7 +2737,7 @@ def _matches_ponyxl_marker(self, text: object) -> bool: haystack = str(text).strip().lower() if not haystack: return False - return any(pattern.search(haystack) for pattern in _LORANADO_MODEL_PATTERNS) + return any(pattern.search(haystack) for pattern in _LORANADO_PONY_PATTERNS) def _is_relevant_pony_metadata_key(self, key: object) -> bool: if key is None: @@ -3807,10 +2874,10 @@ def _prepare_loranado_choice_state( if auto_detect_pony: choice_names = detected_names or all_names if detected_names: - status = f"Detected {len(detected_names)} PonyXL/Anima-compatible LoRAs in `{scan.get('target_folder', '')}`." + status = f"Detected {len(detected_names)} PonyXL-compatible LoRAs in `{scan.get('target_folder', '')}`." elif all_names: status = ( - f"No PonyXL/Anima markers detected in `{scan.get('target_folder', '')}`. " + f"No PonyXL markers detected in `{scan.get('target_folder', '')}`. " f"Falling back to all {len(all_names)} LoRAs." ) else: @@ -3945,7 +3012,7 @@ def _setup_cache(self, use_cache): old_client.close() except Exception as exc: print(f"[R] Warn: Failed to close previous booru session: {exc}") - self._http_client = rb_requesting.BooruSession(use_cache=bool(use_cache)) + self._http_client = rb_http_client.BooruSession(use_cache=bool(use_cache)) print(f"[R] Booru request cache {'enabled' if use_cache else 'disabled'} for this run.") return False @@ -4086,6 +3153,18 @@ def _prepare_tags( def _get_booru_api( self, booru_name, fringe_benefits, gelbooru_credentials: Optional[Dict[str, str]] = None ): + from ranboorux.boorus.gelbooru import Gelbooru, GelbooruCompatible + from ranboorux.boorus.simple import ( + AIBooru, + Danbooru, + Konachan, + Rule34, + Safebooru, + XBooru, + Yandere, + e621, + ) + booru_name = (booru_name or "").strip().lower() if booru_name == "gelbooru-compatible": base_url = _sanitize_gelbooru_compat_base_url( @@ -4230,11 +3309,11 @@ def _fetch_images(self, posts_to_fetch, use_same_image, booru_name, fringe_benef img_to_append = None try: if img_url and img_url.startswith(("http://", "https://")): - safe_url = rb_requesting.redact_url(img_url) + safe_url = rb_http_client.redact_url(img_url) print(f"[R] Fetching {i+1}/{len(image_urls)}: {safe_url[:80]}...") content = self._http_client.get_bytes( img_url, - headers=api.headers, + headers=self._get_image_fetch_headers(api, img_url), timeout=30, max_bytes=MAX_SOURCE_IMAGE_BYTES, ) @@ -4255,16 +3334,16 @@ def _fetch_images(self, posts_to_fetch, use_same_image, booru_name, fringe_benef for site in ["pixiv.net", "pximg.net", "twitter.com", "x.com"] ): print( - f"[R] Skipped external site URL {i+1}: {rb_requesting.redact_url(img_url)[:80]} (not a direct image)" + f"[R] Skipped external site URL {i+1}: {rb_http_client.redact_url(img_url)[:80]} (not a direct image)" ) else: print( - f"[R] Invalid URL protocol {i+1}: {rb_requesting.redact_url(img_url)[:80]}" + f"[R] Invalid URL protocol {i+1}: {rb_http_client.redact_url(img_url)[:80]}" ) else: print(f"[R] No URL available for image {i+1}") except Exception as e: - safe_msg = rb_requesting.safe_exception_message("Image fetch", img_url, e) + safe_msg = rb_http_client.safe_exception_message("Image fetch", img_url, e) print(f"[R] Error fetching image {i+1}: {safe_msg}") fetched_images.append(img_to_append) print(f"[R] Fetched {fetched_count} images.") @@ -4272,6 +3351,19 @@ def _fetch_images(self, posts_to_fetch, use_same_image, booru_name, fringe_benef print("[R] Warn: Some images failed.") return fetched_images + def _get_image_fetch_headers(self, api, img_url: str) -> dict: + base = dict(api.headers) + if "gelbooru" in img_url.lower() or "img4.gelbooru.com" in img_url.lower(): + base.update( + { + "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", + "referer": "https://gelbooru.com/", + "accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8", + "accept-language": "en-US,en;q=0.9", + } + ) + return base + def _process_single_prompt( self, index, raw_prompt, base_positive, base_negative, initial_additions, settings ): @@ -4366,6 +3458,7 @@ def _process_single_prompt( favorites_guard: Set[str] = set() if filter_ctx: favorites_guard = set(filter_ctx.get("favorites", frozenset())) # type: ignore[arg-type] + catalog = self._active_catalog() norm_cache = getattr(self, "_tag_normal_cache", {}) if not isinstance(norm_cache, dict): norm_cache = {} @@ -4394,15 +3487,31 @@ def _process_single_prompt( or (t_norm and (t_norm.endswith(" series") or t_norm.endswith(" franchise"))) ): should_remove = True - if not should_remove and remove_clothing_tags and self._is_clothing_tag(t): + if ( + not should_remove + and remove_clothing_tags + and rb_tag_pipeline.is_clothing_tag(t) + ): should_remove = True - if not should_remove and remove_text_tags and self._is_textual_tag(t): + if ( + not should_remove + and remove_text_tags + and rb_tag_pipeline.is_textual_tag(t, catalog.is_textual if catalog else None) + ): should_remove = True - if not should_remove and remove_furry_tags and self._is_furry_tag(t): + if not should_remove and remove_furry_tags and rb_tag_pipeline.is_furry_tag(t): should_remove = True - if not should_remove and remove_headwear_tags and self._is_headwear_tag(t): + if ( + not should_remove + and remove_headwear_tags + and rb_tag_pipeline.is_headwear_tag(t) + ): should_remove = True - if not should_remove and remove_series_tags and self._is_series_tag(t): + if ( + not should_remove + and remove_series_tags + and rb_tag_pipeline.is_series_tag(t, catalog.category if catalog else None) + ): should_remove = True if not should_remove and preserve_hair_eye_colors: if base_hair_colors and canonical_tag in base_hair_colors: @@ -4411,17 +3520,23 @@ def _process_single_prompt( pass elif ( base_hair_colors - and self._is_hair_color_tag(t) + and rb_tag_pipeline.is_hair_color_tag( + t, catalog.is_hair if catalog else None + ) and canonical_tag not in base_hair_colors ): should_remove = True elif ( base_eye_colors - and self._is_eye_color_tag(t) + and rb_tag_pipeline.is_eye_color_tag(t, catalog.is_eye if catalog else None) and canonical_tag not in base_eye_colors ): should_remove = True - if not should_remove and restrict_subject_tags and self._is_subject_tag(t): + if ( + not should_remove + and restrict_subject_tags + and rb_tag_pipeline.is_subject_tag(t) + ): subject_norm = t_norm if allowed_subjects: if subject_norm not in allowed_subjects: @@ -4452,9 +3567,13 @@ def _process_single_prompt( elif chaos_mode == "Shuffle Negative": _, current_negative = generate_chaos("", current_negative, chaos_amount) if limit_tags_pct < 1.0: - current_prompt = limit_prompt_tags(current_prompt, limit_tags_pct, "Limit") + current_prompt = rb_tag_pipeline.limit_prompt_tags( + current_prompt, limit_tags_pct, "Limit" + ) if max_tags_count > 0: - current_prompt = limit_prompt_tags(current_prompt, max_tags_count, "Max") + current_prompt = rb_tag_pipeline.limit_prompt_tags( + current_prompt, max_tags_count, "Max" + ) if change_dash: current_prompt = current_prompt.replace("_", " ") current_negative = current_negative.replace("_", " ") @@ -4462,8 +3581,8 @@ def _process_single_prompt( current_prompt = ( f"{base_positive}, {current_prompt}" if current_prompt else base_positive ) - current_prompt = remove_repeated_tags(current_prompt) - current_negative = remove_repeated_tags(current_negative) + current_prompt = rb_tag_pipeline.remove_repeated_tags(current_prompt) + current_negative = rb_tag_pipeline.remove_repeated_tags(current_negative) return current_prompt, current_negative def _apply_loranado( @@ -4503,11 +3622,11 @@ def _apply_loranado( detected_loras = list(scan.get("detected_files") or []) if detected_loras: candidate_loras = detected_loras - print(f"[R] LoRAnado: using {len(candidate_loras)} PonyXL/Anima-detected LoRAs.") + print(f"[R] LoRAnado: using {len(candidate_loras)} PonyXL-detected LoRAs.") else: candidate_loras = all_loras print( - f"[R] LoRAnado: no PonyXL/Anima markers detected in {target_folder}; falling back to all LoRAs." + f"[R] LoRAnado: no PonyXL markers detected in {target_folder}; falling back to all LoRAs." ) else: candidate_loras = all_loras @@ -4601,6 +3720,20 @@ def _prepare_img2img_pass(self, p, use_img2img, use_ip): 0.6, self.img2img_denoising ) # Cap at 0.6 to prevent distortion + # Anima-specific img2img overrides + options = getattr(self, "options", None) + if getattr(self, "_is_anima_model", False) and getattr( + options, "anima_tune_img2img", getattr(options, "anima_auto_detect", True) + ): + self.img2img_denoising = min(0.5, self.img2img_denoising) + initial_steps = max(8, min(15, p.steps // 3)) + self._host_scope.set_attr(p, "steps", initial_steps) + p.cfg_scale = max(3.0, min(p.cfg_scale, 6.0)) + print( + f"[R] Anima: using flow-matching optimized parameters " + f"(denoise={self.img2img_denoising}, steps={initial_steps}, cfg={p.cfg_scale})" + ) + self.run_img2img_pass = True self._img2img_final_outpath_samples = getattr(p, "outpath_samples", None) @@ -4649,7 +3782,7 @@ def _cleanup_after_run(self, use_cache): ) except Exception as exc: print(f"[R Cleanup] Host mutation restore failed: {exc}") - self._host_scope = rb_host_state.HostMutationScope() + self._host_scope = rb_mutation_scope.HostMutationScope() # Clean up stored original values if hasattr(self, "original_full_prompt"): @@ -4708,7 +3841,7 @@ def _cleanup_after_run(self, use_cache): http_client.close() except Exception as exc: print(f"[R Post] Warn: Failed to close booru session: {exc}") - self._http_client = rb_requesting.BooruSession(use_cache=False) + self._http_client = rb_http_client.BooruSession(use_cache=False) if hasattr(self, "cache_installed_by_us"): try: del self.cache_installed_by_us @@ -4855,6 +3988,33 @@ def _maybe_release_stale_guards(self, new_processing_obj): processing_obj=new_processing_obj, ) + @staticmethod + def _anima_quality_prefix() -> str: + """Return Anima's recommended positive quality prefix.""" + return "masterpiece, best quality, score_7, safe, " + + @staticmethod + def _anima_negative_default() -> str: + """Return Anima's recommended negative prompt.""" + return "worst quality, low quality, score_1, score_2, score_3, artist name, blurry, jpeg artifacts, chromatic aberration" + + @staticmethod + def _has_quality_prefix(prompt: str) -> bool: + """Check if prompt already has quality tokens (case-insensitive).""" + if not prompt: + return False + quality_tokens = { + "masterpiece", + "best quality", + "high quality", + "score_7", + "score_8", + "score_9", + "safe", + } + first_10 = [t.strip().lower() for t in prompt.split(",")[:10]] + return any(token in tag for token in quality_tokens for tag in first_10) + def before_process(self, p: StableDiffusionProcessing, *args): try: # Fast-path for our own internal img2img calls: initialize seeds and exit @@ -4884,14 +4044,6 @@ def before_process(self, p: StableDiffusionProcessing, *args): print(f"[R Before] WARN: Internal img2img seed init failed: {_e}") return - # Parse the enabled flag EARLY so postprocess always has the current state, - # even if we hit a guard or exception before the full arg-unpack below. - try: - early_options = rb_run_options.RunOptions.from_script_args(args) - self._post_enabled = bool(early_options.enabled) - except Exception: - pass # will be handled by the main unpack later - # Ensure leftover guards from interrupted jobs don't block new generations self._maybe_release_stale_guards(p) @@ -4936,7 +4088,7 @@ def before_process(self, p: StableDiffusionProcessing, *args): self._host_scope.restore() except Exception as exc: print(f"[R Before] Warn: stale host-scope cleanup failed: {exc}") - self._host_scope = rb_host_state.HostMutationScope() + self._host_scope = rb_mutation_scope.HostMutationScope() script_args_source = getattr(p, "script_args", None) if isinstance(script_args_source, (list, tuple)): self._adetailer_script_args_snapshot = list(script_args_source) @@ -5038,6 +4190,20 @@ def before_process(self, p: StableDiffusionProcessing, *args): self._manual_adetailer_prev_enabled = self._adetailer_support_enabled self._log_prompt_sources = bool(log_prompt_sources_ui) + # Anima model detection + try: + info = get_anima_model_info(shared.sd_model) + self._is_anima_model = info["detected"] + if self._is_anima_model: + anima_auto_detect = getattr(options, "anima_auto_detect", True) + if anima_auto_detect: + change_dash = True + print( + f"[R] Anima model detected ({info['model_name']}) - auto-enabling space-separated tags" + ) + except Exception: + self._is_anima_model = False + self._current_booru_name = booru if booru == "gelbooru": self._gelbooru_effective_credentials = self._resolve_gelbooru_credentials( @@ -5203,6 +4369,15 @@ def before_process(self, p: StableDiffusionProcessing, *args): if isinstance(p.prompt, str) else (p.prompt[0] if isinstance(p.prompt, list) and p.prompt else "") ) + + # Anima: apply default prompts + if self._is_anima_model and getattr(options, "anima_auto_detect", True): + if not self._has_quality_prefix(self.original_prompt): + self.original_prompt = f"{self._anima_quality_prefix()}{self.original_prompt}" + print("[R] Anima: applied default quality tags") + if not isinstance(p.negative_prompt, str) or not p.negative_prompt.strip(): + p.negative_prompt = self._anima_negative_default() + base_hair_colors, base_eye_colors = self._extract_color_tags(self.original_prompt) self._base_hair_color_tags = base_hair_colors self._base_eye_color_tags = base_eye_colors @@ -5512,69 +4687,58 @@ def before_process(self, p: StableDiffusionProcessing, *args): if use_ip and self.last_img and self.last_img[0] is not None: cn_configured = False - # Preferred: external_code API from ControlNet + # Forge Neo direct: find ControlNet script in alwayson_scripts try: - cn_module = self._load_cn_external_code() - if hasattr(cn_module, "get_all_units_in_processing") and hasattr( - cn_module, "update_cn_script_in_processing" - ): - cn_units = cn_module.get_all_units_in_processing(p) - if cn_units and len(cn_units) > 0: - copied_unit = cn_units[0].__dict__.copy() - copied_unit["enabled"] = True - copied_unit["weight"] = float(self.img2img_denoising) - img_for_cn = ( - self.last_img[0].convert("RGB") - if self.last_img[0].mode != "RGB" - else self.last_img[0] + scripts_runner = getattr(p, "scripts", None) + cn_script = None + if scripts_runner is not None: + for s in getattr(scripts_runner, "alwayson_scripts", []): + filename = getattr(s, "filename", "") or "" + title = getattr(s, "title", lambda: "")() + if "controlnet" in filename.lower() or "controlnet" in title.lower(): + cn_script = s + break + + if cn_script is not None: + start = getattr(cn_script, "args_from", None) + end = getattr(cn_script, "args_to", None) + if isinstance(start, int) and isinstance(end, int) and 0 <= start < end: + full_args = ( + list(p.script_args) + if isinstance(p.script_args, tuple) + else list(p.script_args or []) ) - copied_unit["image"]["image"] = np.array(img_for_cn) - cn_module.update_cn_script_in_processing( - p, [copied_unit] + cn_units[1:] - ) - cn_configured = True - print("[R Before] ControlNet configured via external_code.") - # else: module loaded but does not expose update helpers; silently skip to fallback - except Exception: - # Silently fallback if external_code path not supported in this build - pass - - # Fallback: p.script_args hack (fragile but effective) - if not cn_configured: - cn_arg_start_guess = 0 - num_controls_per_unit = 20 - if num_controls_per_unit > 0: - target_unit_arg_start = cn_arg_start_guess - enabled_idx = target_unit_arg_start + 0 - weight_idx = target_unit_arg_start + 3 - image_idx = target_unit_arg_start + 4 - args_source = p.script_args - if isinstance(args_source, tuple): - args_target_list = list(args_source) - max_idx = max(enabled_idx, weight_idx, image_idx) - if max_idx < len(args_target_list): - try: - img_for_cn = ( - self.last_img[0].convert("RGB") - if self.last_img[0].mode != "RGB" - else self.last_img[0] - ) - cn_image_input = {"image": np.array(img_for_cn), "mask": None} - args_target_list[enabled_idx] = True - args_target_list[weight_idx] = float(self.img2img_denoising) - args_target_list[image_idx] = cn_image_input - p.script_args = tuple(args_target_list) - print( - "[R Before] ControlNet using fallback p.script_args hack." - ) - except Exception as e: - print(f"[R Before] Error setting CN via p.script_args: {e}") - else: + if end <= len(full_args): + unit = full_args[start] + img_for_cn = ( + self.last_img[0].convert("RGB") + if self.last_img[0].mode != "RGB" + else self.last_img[0] + ) + cn_image = {"image": np.array(img_for_cn), "mask": None} + + if isinstance(unit, dict): + unit["enabled"] = True + unit["weight"] = float(self.img2img_denoising) + unit["image"] = cn_image + elif hasattr(unit, "enabled"): + unit.enabled = True + unit.weight = float(self.img2img_denoising) + unit.image = cn_image + + setattr(p, "resize_mode", 1) + p.script_args = tuple(full_args) + cn_configured = True print( - f"[R Before] Error: CN arg index ({max_idx}) OOB ({len(args_target_list)})." + "[R Before] ControlNet configured via Forge Neo direct (p.script_args slice)." ) - else: - print("[R Before] Error: p.script_args is not a tuple.") + except Exception as e: + print(f"[R Before] ControlNet config error: {e}") + + if not cn_configured and use_ip: + if not hasattr(p, "resize_mode"): + setattr(p, "resize_mode", 1) + print("[R Before] ControlNet script not found; p.resize_mode safeguard set.") self._prepare_img2img_pass(p, use_img2img, use_ip) @@ -5635,259 +4799,14 @@ def _reset_adetailer_state_for_run(self, p): def _restore_native_adetailer_scripts(self, p): """Ensure native ADetailer scripts resume running when manual support is disabled.""" - try: - if not self._adetailer_patches.is_empty(): - self._unpatch_manual_adetailer_overrides() - except Exception as exc: - print(f"[R Before] Warn: Could not unpatch manual ADetailer overrides: {exc}") - try: - self._set_adetailer_block(False) - except Exception: - pass - setattr(self.__class__, "_ranbooru_block_all_adetailer", False) - setattr(self.__class__, "_adetailer_global_guard_active", False) - try: - self._restore_early_adetailer_protection(p) - except Exception as exc: - print(f"[R Before] Warn: Could not restore ADetailer runner state: {exc}") - try: - self._reenable_adetailer_from_previous_generation() - except Exception as exc: - print(f"[R Before] Warn: Could not re-enable ADetailer scripts: {exc}") - try: - restored = self._force_enable_adetailer_scripts(p) - except Exception as exc: - print(f"[R Before] Warn: Could not force-enable ADetailer scripts: {exc}") - restored = 0 - if restored: - print( - f"[R Before] Restored {restored} native ADetailer script(s) after manual toggle was disabled" - ) - if hasattr(self, "disabled_adetailer_scripts"): - try: - delattr(self, "disabled_adetailer_scripts") - except Exception: - pass - guard_present = False - try: - import modules.scripts as scripts_module - - for runner_attr in ("scripts_txt2img", "scripts_img2img"): - runner = getattr(scripts_module, runner_attr, None) - if runner and getattr(runner, "_ranbooru_guard_installed", False): - guard_present = True - break - except Exception: - guard_present = False - if guard_present: - try: - self._reset_script_runner_guards() - except Exception as exc: - print(f"[R Before] Warn: Could not reset script runner guards: {exc}") - self._ensure_native_adetailer_enable_flags(p) - if not self._native_adetailer_detected(): - try: - import modules.scripts as scripts_module - - if hasattr(scripts_module, "reload_scripts"): - print("[R Before] Reloading scripts to restore native ADetailer") - scripts_module.reload_scripts() - except Exception as exc: - print(f"[R Before] Warn: Could not reload scripts for ADetailer: {exc}") + self._adetailer_orch._restore_native_adetailer_scripts(p) def _force_enable_adetailer_scripts(self, processing_obj=None): """Return the count of ADetailer scripts restored to their original behaviour.""" - try: - import modules.scripts as scripts_module - except Exception as exc: - print(f"[R Before] Warn: Could not access scripts module to restore ADetailer: {exc}") - return 0 - runners = [] - for runner_attr in ("scripts_txt2img", "scripts_img2img"): - runner = getattr(scripts_module, runner_attr, None) - if runner: - runners.append(runner) - if ( - processing_obj is not None - and hasattr(processing_obj, "scripts") - and processing_obj.scripts not in runners - ): - runners.append(processing_obj.scripts) - seen_ids = set() - restored_count = 0 - for runner in runners: - if runner is None: - continue - for list_attr in ("alwayson_scripts", "scripts"): - script_list = getattr(runner, list_attr, None) - if not script_list: - continue - for script in script_list: - if not script: - continue - script_id = id(script) - if script_id in seen_ids: - continue - seen_ids.add(script_id) - if not self._is_adetailer_script(script): - continue - restored = False - if hasattr(script, "enabled") and script.enabled is False: - script.enabled = True - restored = True - for method_name in ( - "postprocess", - "process", - "process_batch", - "before_process", - "after_process", - ): - backup_name = f"_ranbooru_original_{method_name}" - if hasattr(script, backup_name): - try: - setattr(script, method_name, getattr(script, backup_name)) - except Exception: - pass - try: - delattr(script, backup_name) - except Exception: - pass - restored = True - for attr in ("_ranbooru_disabled_after_manual", "_ranbooru_disabled_source"): - if hasattr(script, attr): - try: - delattr(script, attr) - except Exception: - pass - restored = True - if restored: - restored_count += 1 - if restored_count == 0: - try: - debug_entries = [] - for runner in runners: - if not runner: - continue - for list_attr in ("alwayson_scripts", "scripts"): - script_list = getattr(runner, list_attr, None) - if not script_list: - continue - for script in script_list: - if self._is_adetailer_script(script): - debug_entries.append( - f"{script.__class__.__name__}(enabled={getattr(script, 'enabled', 'n/a')})" - ) - if debug_entries: - print( - f"[R Before] Native ADetailer scripts detected: {', '.join(debug_entries)}" - ) - except Exception: - pass - return restored_count + return self._adetailer_orch._force_enable_adetailer_scripts(processing_obj) def _ensure_native_adetailer_enable_flags(self, processing_obj): - if not getattr(self, "_adetailer_support_enabled", False): - return - try: - args = getattr(processing_obj, "script_args", None) - except Exception as exc: - print(f"[R Before] Native ADetailer: unable to read script_args: {exc}") - return - if not isinstance(args, (list, tuple)) or not args: - print( - "[R Before] Native ADetailer: script_args empty or not list/tuple; skipping flag repair" - ) - return - args_list = list(args) - runners = [] - runner = getattr(processing_obj, "scripts", None) - if runner is not None: - runners.append(runner) - try: - import modules.scripts as scripts_module - - for attr in ("scripts_txt2img", "scripts_img2img"): - global_runner = getattr(scripts_module, attr, None) - if global_runner is not None and global_runner not in runners: - runners.append(global_runner) - except Exception as exc: - print(f"[R Before] Native ADetailer: could not gather global runners: {exc}") - candidates = [] - for r in runners: - for list_attr in ("alwayson_scripts", "scripts"): - script_list = getattr(r, list_attr, None) - if script_list: - candidates.extend(script_list) - if not candidates: - print("[R Before] Native ADetailer: no script candidates found for flag repair") - return - changed = False - for script in candidates: - if not self._is_adetailer_script(script): - continue - extracted = self._extract_adetailer_script_args(script, processing_obj) - sanitized = list(extracted.get("args") or []) - meta = extracted.get("meta") or {} - start_idx = meta.get("slice_start") - end_idx = meta.get("slice_end") - if start_idx is None or end_idx is None: - continue - start_idx = max(0, min(len(args_list), start_idx)) - end_idx = max(start_idx, min(len(args_list), end_idx)) - if not sanitized or end_idx - start_idx != len(sanitized): - slice_view = args_list[start_idx:end_idx] - else: - slice_view = sanitized - print( - f"[R Before] Native ADetailer candidate {script.__class__.__name__} enabled={getattr(script, 'enabled', 'n/a')} slice [{start_idx}:{end_idx}] -> {slice_view}" - ) - if not sanitized: - continue - bool_index = 0 - local_changed = False - for offset, val in enumerate(sanitized): - if isinstance(val, bool): - if bool_index == 0 and val is False: - sanitized[offset] = True - local_changed = True - print( - f"[R Before] Set native ADetailer enable flag True at offset {offset}" - ) - elif bool_index == 1 and val is True: - sanitized[offset] = False - local_changed = True - print(f"[R Before] Cleared native ADetailer skip flag at offset {offset}") - bool_index += 1 - elif isinstance(val, dict): - if val.get("ad_tab_enable") is False and val.get("ad_model") not in ( - None, - "", - "None", - ): - val["ad_tab_enable"] = True - local_changed = True - print(f"[R Before] Enabled ad_tab_enable in dict at offset {offset}") - if local_changed: - if end_idx - start_idx == len(sanitized): - args_list[start_idx:end_idx] = sanitized - changed = True - continue - # fallback if lengths mismatch - for offset, val in enumerate(sanitized): - target_idx = start_idx + offset - if target_idx < len(args_list): - args_list[target_idx] = val - else: - args_list.append(val) - changed = True - if changed: - if isinstance(args, list): - processing_obj.script_args = args_list - else: - processing_obj.script_args = tuple(args_list) - print(f"[R Before] Native ADetailer flags updated: {args_list}") - else: - print("[R Before] Native ADetailer flags already enabled; no changes made") + self._adetailer_orch._ensure_native_adetailer_enable_flags(processing_obj) def _force_native_adetailer_execution(self, p, processed): if getattr(self, "_adetailer_support_enabled", False): @@ -6016,14 +4935,7 @@ def postprocess(self, p: StableDiffusionProcessing, processed, *args): "[R Post] Img2Img already started for this generation; skipping duplicate postprocess entry" ) return - # Read enabled from script args first (authoritative), fall back to instance var enabled = getattr(self, "_post_enabled", False) - try: - if args: - post_options = rb_run_options.RunOptions.from_script_args(args) - enabled = bool(post_options.enabled) - except Exception: - pass # keep the _post_enabled fallback use_img2img = getattr(self, "_post_use_img2img", False) getattr(self, "_post_use_last_img", False) crop_center = getattr(self, "_post_crop_center", False) @@ -6114,7 +5026,7 @@ def postprocess(self, p: StableDiffusionProcessing, processed, *args): f"[R Post] Preparing {len(self.last_img)} images ({'Crop' if crop_center else 'Resize'}) to {target_w}x{target_h} for Img2Img." ) prepared_images = [ - resize_image(img, target_w, target_h, cropping=crop_center) + rb_image_ops.resize_image(img, target_w, target_h, cropping=crop_center) for img in self.last_img if img is not None ] @@ -6359,7 +5271,7 @@ def postprocess(self, p: StableDiffusionProcessing, processed, *args): self._clear_processing_guards(p) def _is_adetailer_enabled(self): - return getattr(self, "_adetailer_support_enabled", False) + return self._adetailer_orch.is_adetailer_enabled() def _set_adetailer_block(self, should_block: bool): """Toggle the global guard on patched ADetailer classes""" @@ -6541,65 +5453,7 @@ def _images_visibly_different(self, original_image, processed_image): def _execute_manual_adetailer(self, p, processed, img2img_results): """Run manual ADetailer on img2img results via the deterministic runtime executor.""" - if not self._is_adetailer_enabled() or not img2img_results: - return False - - self._clear_manual_adetailer_skip_flags(p) - adetailer_scripts = rb_adetailer_runtime.gather_adetailer_scripts(p) - if not adetailer_scripts: - print("[R Post] WARN: No ADetailer scripts discovered for manual execution") - return False - - setattr(self.__class__, "_ranbooru_manual_adetailer_active", True) - - def build_processed(single_image): - temp_processed = types.SimpleNamespace() - temp_processed.images = [single_image] - temp_processed.image = single_image - for attr in ( - "prompt", - "negative_prompt", - "seed", - "subseed", - "width", - "height", - "cfg_scale", - "steps", - ): - if hasattr(processed, attr): - setattr(temp_processed, attr, getattr(processed, attr)) - return temp_processed - - try: - result = rb_adetailer_runtime.execute_manual_adetailer( - adetailer_scripts=adetailer_scripts, - images=list(img2img_results), - processing_obj=p, - run_state=self._adetailer_state, - patch_registry=self._adetailer_patches, - extract_script_args=self._extract_adetailer_script_args, - build_processed=build_processed, - isolation_factory=lambda script_obj: self._manual_adetailer_script_isolation( - p, - script_obj, - keep_controlnet=self._manual_adetailer_requires_controlnet( - self._extract_adetailer_script_args(script_obj, p).get("args") or [] - ), - ), - ) - finally: - setattr(self.__class__, "_ranbooru_manual_adetailer_active", False) - - for error in result.errors: - print(f"[R Post] WARN: Manual ADetailer error: {error}") - processed.images.clear() - processed.images.extend(result.images) - img2img_results.clear() - img2img_results.extend(result.images) - if hasattr(p, "processed") and hasattr(p.processed, "images"): - p.processed.images.clear() - p.processed.images.extend(result.images) - return result.successful_processes > 0 + return self._adetailer_orch._execute_manual_adetailer(p, processed, img2img_results) def _unpatch_manual_adetailer_overrides(self): """Restore any monkey patches applied for manual ADetailer runs.""" @@ -6640,26 +5494,7 @@ def _unpatch_manual_adetailer_overrides(self): def _is_adetailer_script(self, script): """Check if a script is an ADetailer script""" - try: - if script is None: - return False - script_name = ( - script.__class__.__name__.lower() - if hasattr(script, "__class__") - else str(script).lower() - ) - return ( - "adetailer" in script_name - or "afterdetailer" in script_name - or "after_detailer" in script_name - or "ad_script" in script_name - ) - except Exception as exc: - _ranbooru_logger.warning( - "Failed to inspect ADetailer script type: %s", - rb_requesting.sanitize_exception_text(str(exc)), - ) - return False + return self._adetailer_orch._is_adetailer_script(script) def _is_controlnet_script(self, script): """Check if a script appears to be a ControlNet script.""" @@ -6766,83 +5601,11 @@ def keep_controlnet_fn(script_item, list_attr): def _mark_initial_pass(self, p): """Mark that we're in initial pass so ADetailer can be intercepted later""" - try: - print("[R] Marking initial pass - ADetailer will run on img2img results instead") - - # Clear any previous hard-disable flag for ADetailer - try: - if hasattr(p, "_ad_disabled") and getattr(p, "_ad_disabled", False): - self._host_scope.set_attr(p, "_ad_disabled", False) - print("[R] Cleared p._ad_disabled from previous generation") - except Exception as _e: - print(f"[R] WARN: Could not clear p._ad_disabled: {_e}") - - # Clear our class-level guard - self._set_adetailer_block(False) - self._adetailer_state.initial_pass_suppressed = False - # Clear pipeline-level guard flag - setattr(self.__class__, "_ranbooru_block_all_adetailer", False) - - # Install runner guard (idempotent) - self._install_scriptrunner_guard(p) - - # CRITICAL: Re-enable any ADetailer scripts from previous generation - self._reenable_adetailer_from_previous_generation() - - # Just set a flag that we're in initial pass - self._ranbooru_initial_pass = True - - # Store reference to processing object for later use - self._initial_pass_p = p - - except Exception as e: - print(f"[R] Error marking initial pass: {e}") + self._adetailer_orch._mark_initial_pass(p) def _reenable_adetailer_from_previous_generation(self): """Re-enable ALL ADetailer scripts that were disabled in the previous generation""" - try: - if hasattr(self, "disabled_adetailer_scripts") and self.disabled_adetailer_scripts: - print( - f"[R] COMPREHENSIVE RE-ENABLE: Restoring {len(self.disabled_adetailer_scripts)} ADetailer script(s) from previous generation" - ) - - for script, original_enabled in self.disabled_adetailer_scripts: - source = getattr(script, "_ranbooru_disabled_source", "unknown") - print(f"[R] Re-enabling {script.__class__.__name__} from {source}") - - # Restore original enabled state - if hasattr(script, "enabled"): - script.enabled = original_enabled - - # Restore ALL original methods that were disabled - methods_to_restore = [ - "postprocess", - "process", - "process_batch", - "before_process", - "after_process", - ] - for method_name in methods_to_restore: - original_method_attr = f"_ranbooru_original_{method_name}" - if hasattr(script, original_method_attr): - original_method = getattr(script, original_method_attr) - setattr(script, method_name, original_method) - delattr(script, original_method_attr) - - # Remove our disable flags - if hasattr(script, "_ranbooru_disabled_after_manual"): - delattr(script, "_ranbooru_disabled_after_manual") - if hasattr(script, "_ranbooru_disabled_source"): - delattr(script, "_ranbooru_disabled_source") - - print( - f"[R] COMPREHENSIVE RE-ENABLE: Restored {len(self.disabled_adetailer_scripts)} ADetailer script(s) for new generation" - ) - # Clear the list now that we've re-enabled everything - delattr(self, "disabled_adetailer_scripts") - - except Exception as e: - print(f"[R] Error in comprehensive ADetailer re-enable: {e}") + self._adetailer_orch._reenable_adetailer_from_previous_generation() def _prevent_all_image_saving(self, p, temp_dir): """Prevent all possible image saving during initial pass""" @@ -6882,16 +5645,7 @@ def _prevent_all_image_saving(self, p, temp_dir): def _prepare_adetailer_for_img2img(self, p): """Prepare ADetailer to run on img2img results""" - if not self._is_adetailer_enabled(): - return - try: - print("[R] Preparing ADetailer to run on img2img results") - - # Clear the initial pass flag so ADetailer knows to run normally - self._ranbooru_initial_pass = False - - except Exception as e: - print(f"[R] Error preparing ADetailer: {e}") + self._adetailer_orch._prepare_adetailer_for_img2img(p) def _force_ui_update(self, p, processed, final_results): """Force ForgeUI to display our final ADetailer-processed results""" @@ -6970,7 +5724,7 @@ def _force_ui_update(self, p, processed, final_results): except Exception as ui_update_error: _ranbooru_logger.warning( "Unable to add UI force-update flags: %s", - rb_requesting.sanitize_exception_text(str(ui_update_error)), + rb_http_client.sanitize_exception_text(str(ui_update_error)), ) # Method 5: Update the main result that ForgeUI looks for @@ -7039,127 +5793,15 @@ def process(self, p, *args): def _early_adetailer_protection(self, p): """Complete ADetailer blocking during initial pass - remove scripts entirely""" - if not self._is_adetailer_enabled(): - return - try: - print("[R Process] Early ADetailer protection activated") - - # Check if we're in the initial pass - if getattr(self, "_ranbooru_initial_pass", False): - print("[R Process] Detected initial pass - COMPLETELY BLOCKING ADetailer") - - # Set comprehensive block flags - self._host_scope.set_attr(p, "_ranbooru_skip_initial_adetailer", True) - self._host_scope.set_attr(p, "_ranbooru_suppress_all_processing", True) - self._host_scope.set_attr(p, "_ranbooru_initial_pass_only", True) - self._host_scope.set_attr(p, "_ad_disabled", True) - self._adetailer_state.initial_pass_suppressed = True - - # CRITICAL: Completely remove ADetailer scripts from the runner during initial pass - self._remove_adetailer_from_runner(p) - - # Set multiple block flags to ensure no ADetailer execution - self._set_adetailer_block(True) - setattr(self.__class__, "_ranbooru_block_all_adetailer", True) - setattr(self.__class__, "_adetailer_global_guard_active", True) - self._adetailer_state.global_guard_active = True - - print( - "[R Process] ADetailer completely blocked for initial pass - will be restored for manual img2img processing" - ) - - except Exception as e: - print(f"[R Process] Error in early ADetailer protection: {e}") + self._adetailer_orch._early_adetailer_protection(p) def _remove_adetailer_from_runner(self, p): """Temporarily remove ADetailer scripts from the script runner during initial pass""" - try: - if not hasattr(p, "scripts") or p.scripts is None: - return - - # Store original scripts for restoration - if not hasattr(self, "_stored_adetailer_scripts"): - self._stored_adetailer_scripts = {"alwayson": [], "regular": []} - - # Remove ADetailer from alwayson_scripts - if hasattr(p.scripts, "alwayson_scripts") and p.scripts.alwayson_scripts: - original_alwayson = list(p.scripts.alwayson_scripts) - filtered_alwayson = [ - s for s in original_alwayson if not self._is_adetailer_script(s) - ] - removed_alwayson = [s for s in original_alwayson if self._is_adetailer_script(s)] - - p.scripts.alwayson_scripts = filtered_alwayson - self._stored_adetailer_scripts["alwayson"] = removed_alwayson - print( - f"[R Process] Removed {len(removed_alwayson)} ADetailer scripts from alwayson_scripts" - ) - - # Remove ADetailer from regular scripts - if hasattr(p.scripts, "scripts") and p.scripts.scripts: - original_scripts = list(p.scripts.scripts) - filtered_scripts = [s for s in original_scripts if not self._is_adetailer_script(s)] - removed_scripts = [s for s in original_scripts if self._is_adetailer_script(s)] - - p.scripts.scripts = filtered_scripts - self._stored_adetailer_scripts["regular"] = removed_scripts - print(f"[R Process] Removed {len(removed_scripts)} ADetailer scripts from scripts") - - except Exception as e: - print(f"[R Process] Error removing ADetailer from runner: {e}") + self._adetailer_orch._remove_adetailer_from_runner(p) def _restore_early_adetailer_protection(self, processing_obj=None): """Restore ADetailer scripts and flags after an interrupted or completed run.""" - try: - print("[R Process] Restoring ADetailer scripts for manual processing") - - # Clear initial pass/block flags so subsequent generations can run ADetailer - setattr(self.__class__, "_ranbooru_block_all_adetailer", False) - setattr(self.__class__, "_adetailer_global_guard_active", False) - self._set_adetailer_block(False) - - # Determine which processing object's script runner to restore into - candidate_p = ( - processing_obj - or getattr(self, "_initial_pass_p", None) - or getattr(self, "_current_processing_object", None) - ) - runner = getattr(candidate_p, "scripts", None) if candidate_p else None - - # Restore scripts we removed during the initial pass safeguard - stored = getattr(self, "_stored_adetailer_scripts", None) - if stored and runner: - try: - if hasattr(runner, "alwayson_scripts") and stored.get("alwayson"): - for script in stored["alwayson"]: - if script not in runner.alwayson_scripts: - runner.alwayson_scripts.append(script) - print( - f"[R Process] Reattached {len(stored['alwayson'])} ADetailer always-on script(s)" - ) - if hasattr(runner, "scripts") and stored.get("regular"): - for script in stored["regular"]: - if script not in runner.scripts: - runner.scripts.append(script) - print( - f"[R Process] Reattached {len(stored['regular'])} ADetailer on-demand script(s)" - ) - finally: - # Clear stored references so we don't duplicate reinsertion - delattr(self, "_stored_adetailer_scripts") - - # Ensure any scripts we hard-disabled are re-enabled for the next generation - if hasattr(self, "disabled_adetailer_scripts"): - self._reenable_adetailer_from_previous_generation() - - # Clear temporary protection flag if present - if hasattr(self, "_temp_disabled_adetailer"): - delattr(self, "_temp_disabled_adetailer") - - print("[R Process] Early protection restoration complete") - - except Exception as e: - print(f"[R Process] Error restoring early ADetailer protection: {e}") + self._adetailer_orch._restore_early_adetailer_protection(processing_obj) def process_batch_pre(self, p, *args, **kwargs): """Pre-batch processing to set up result interception""" @@ -7204,27 +5846,7 @@ def use_autotagger(self, model): def _install_scriptrunner_guard(self, p): """Wrap p.scripts postprocess and postprocess_image to skip ADetailer when our block flag is active""" - try: - if not hasattr(p, "scripts") or p.scripts is None: - return - runner = p.scripts - if getattr(runner, "_ranbooru_guard_installed", False): - return - rb_adetailer_runtime.install_runner_guard( - runner=runner, - block_flag_fn=lambda: bool( - getattr(self.__class__, "_ranbooru_block_all_adetailer", False) - and not getattr(self.__class__, "_ranbooru_manual_adetailer_active", False) - ), - patch_registry=self._adetailer_patches, - ) - runner._ranbooru_guard_installed = True - self._log_patch_event( - "info", "Installed ScriptRunner guard to skip ADetailer when blocked" - ) - except Exception as e: - self._log_patch_event("warning", f"Failed to install ScriptRunner guard: {e}") - print(f"[R] Error installing ScriptRunner guard: {e}") + self._adetailer_orch._install_scriptrunner_guard(p) def _prepare_processing_for_manual_adetailer(self, p, processed, img2img_results): """Ensure p has correct images, sizes, prompts, and save paths before running ADetailer manually""" @@ -7301,57 +5923,7 @@ def _clear_manual_adetailer_skip_flags(self, processing_obj): def _install_preview_guard(self): """Install a guard around shared.state.assign_current_image to block wrong previews""" - try: - import modules.shared as shared - - if not hasattr(shared, "state"): - return - state = shared.state - installed_wrapper = getattr(state, "_ranbooru_preview_guard_wrapper", None) - if ( - getattr(state, "_ranbooru_preview_guard_installed", False) - and installed_wrapper is not None - and getattr(state, "assign_current_image", None) is installed_wrapper - ): - return - if not hasattr(state, "assign_current_image"): - return - original_assign_current_image = state.assign_current_image - script_class = self.__class__ - - def guarded_assign_current_image(img): - try: - if getattr(script_class, "_ranbooru_preview_guard_on", False): - if getattr(script_class, "_ranbooru_preview_block_all", False): - if not getattr( - script_class, "_ranbooru_preview_block_notice_emitted", False - ): - print( - "[R UI] Preview blocked: withholding intermediary frame until final image is ready" - ) - script_class._ranbooru_preview_block_notice_emitted = True - return - # If we know final dims, only allow those; otherwise block 640x512 - final_dims = getattr(script_class, "_ranbooru_final_dims", None) - if img is not None and hasattr(img, "size"): - if final_dims and img.size != final_dims: - print("[R UI] Preview blocked: mismatched size") - return - if img.size == (640, 512): - print("[R UI] Preview blocked: 640x512 preview") - return - except Exception: - pass - return original_assign_current_image(img) - - self._host_scope.patch_attr(state, "assign_current_image", guarded_assign_current_image) - self._host_scope.set_attr(state, "_ranbooru_preview_guard_installed", True) - self._host_scope.set_attr( - state, "_ranbooru_preview_guard_wrapper", guarded_assign_current_image - ) - print("[R UI] Installed preview guard") - except Exception as e: - print(f"[R UI] Error installing preview guard: {e}") + self._adetailer_orch._install_preview_guard() def _set_preview_guard(self, enabled: bool, final_dims=None, block_all: bool = False): try: diff --git a/tests/test_adetailer.py b/tests/test_adetailer.py index f8b75dc..deecbd7 100644 --- a/tests/test_adetailer.py +++ b/tests/test_adetailer.py @@ -565,9 +565,9 @@ def __init__(self): def postprocess_image(self, p, temp_processed, *args): self.calls.append(temp_processed.image.token) - temp_processed.images = [ - DummyImage(f"{temp_processed.image.token}-ad", temp_processed.image.size) - ] + temp_processed.image = DummyImage( + f"{temp_processed.image.token}-ad", temp_processed.image.size + ) return True adetailer_script = AfterDetailerScript() @@ -611,7 +611,7 @@ class AfterDetailerScript: def postprocess_image(self, p, temp_processed, *args): if getattr(p, "_ad_disabled", False): return True - temp_processed.images = [DummyImage(f"{temp_processed.image.token}-ad")] + temp_processed.image = DummyImage(f"{temp_processed.image.token}-ad") return True adetailer_script = AfterDetailerScript() @@ -766,7 +766,7 @@ def __init__(self, **kwargs): subseed=getattr(proc, "subseed", 0), ), ) - monkeypatch.setattr(ranbooru, "resize_image", lambda img, *_args, **_kwargs: img) + monkeypatch.setattr(ranbooru.rb_image_ops, "resize_image", lambda img, *_args, **_kwargs: img) monkeypatch.setattr(script, "_force_ui_update", lambda *_args, **_kwargs: None) monkeypatch.setattr( script, "_prepare_processing_for_manual_adetailer", lambda *_args, **_kwargs: None diff --git a/tests/test_adetailer_runtime.py b/tests/test_adetailer_runtime.py index a0c0b15..021a656 100644 --- a/tests/test_adetailer_runtime.py +++ b/tests/test_adetailer_runtime.py @@ -265,7 +265,8 @@ def postprocess(self, *_args, **_kwargs): def test_execute_manual_adetailer_counts_changed_image(): class AfterDetailerScript: def postprocess_image(self, _p, temp_processed, *_args): - temp_processed.images = [DummyImage(f"{temp_processed.image.token}-ad")] + # Real ADetailer updates pp.image, not pp.images + temp_processed.image = DummyImage(f"{temp_processed.image.token}-ad") state = adetailer_runtime.AdetailerRunState() result = adetailer_runtime.execute_manual_adetailer( @@ -287,7 +288,7 @@ def postprocess_image(self, _p, temp_processed, *_args): def test_execute_manual_adetailer_treats_unchanged_as_noop(): class AfterDetailerScript: def postprocess_image(self, _p, temp_processed, *_args): - temp_processed.images = [temp_processed.image] + temp_processed.image = temp_processed.image result = adetailer_runtime.execute_manual_adetailer( adetailer_scripts=[AfterDetailerScript()], @@ -308,9 +309,9 @@ class AfterDetailerScript: def postprocess_image(self, _p, temp_processed, *_args): token = temp_processed.image.token if token == "img-2": - temp_processed.images = [temp_processed.image] + temp_processed.image = temp_processed.image else: - temp_processed.images = [DummyImage(f"{token}-ad")] + temp_processed.image = DummyImage(f"{token}-ad") result = adetailer_runtime.execute_manual_adetailer( adetailer_scripts=[AfterDetailerScript()], @@ -331,7 +332,7 @@ class AfterDetailerScript: def postprocess_image(self, _p, temp_processed, *_args): if temp_processed.image.token == "img-2": raise RuntimeError("simulated failure") - temp_processed.images = [DummyImage(f"{temp_processed.image.token}-ad")] + temp_processed.image = DummyImage(f"{temp_processed.image.token}-ad") state = adetailer_runtime.AdetailerRunState() result = adetailer_runtime.execute_manual_adetailer( diff --git a/tests/test_anima_detect.py b/tests/test_anima_detect.py new file mode 100644 index 0000000..548a28e --- /dev/null +++ b/tests/test_anima_detect.py @@ -0,0 +1,77 @@ +from ranboorux.anima_detect import get_anima_model_info, is_anima_model + + +class _Obj: + """Minimal attribute holder for test mocks.""" + + pass + + +def test_is_anima_model_none(): + assert is_anima_model(None) is False + + +def test_is_anima_model_non_anima(): + obj = _Obj() + obj.sd_model_checkpoint = "sd_xl_base_1.0.safetensors" + assert is_anima_model(obj) is False + + +def test_is_anima_model_filename_detection(): + obj = _Obj() + obj.sd_model_checkpoint = "anima-base-v1.0.safetensors" + assert is_anima_model(obj) is True + + +def test_is_anima_model_class_detection(): + # Class name containing "Anima" -> True (no checkpoint at all) + obj = type("Anima", (), {})() + assert is_anima_model(obj) is True + + +def test_get_anima_model_info_returns_dict(): + info = get_anima_model_info(None) + assert isinstance(info, dict) + assert "detected" in info + assert "method" in info + assert "model_name" in info + + +def test_is_anima_model_case_insensitive(): + obj = _Obj() + obj.sd_model_checkpoint = "Anima-Base-v1.0.safetensors" + assert is_anima_model(obj) is True + + +def test_is_anima_model_multiple_attr_paths(): + # Fallback to 'checkpoint' attr + obj = _Obj() + obj.checkpoint = "anima-preview3-base.safetensors" + assert is_anima_model(obj) is True + + # Fallback to 'model_checkpoint' attr + obj2 = _Obj() + obj2.model_checkpoint = "anima-aesthetic-v1.0.safetensors" + assert is_anima_model(obj2) is True + + +def test_anima_tune_img2img_can_be_disabled(monkeypatch): + import types + + import scripts.ranbooru as ranbooru + from ranboorux.run_options import RunOptions + + script = ranbooru.Script() + script._is_anima_model = True + script.img2img_denoising = 0.8 + + p = types.SimpleNamespace( + prompt="test", steps=30, cfg_scale=7.5, outpath_samples=None, batch_size=1 + ) + + # When anima_tune_img2img is False, script.img2img_denoising and p.steps should not be overridden by Anima bounds + opts = RunOptions.from_script_args([object()] * 63 + [False]) + script.options = opts + script._prepare_img2img_pass(p, use_img2img=True, use_ip=False) + + assert script.img2img_denoising == 0.6 # Default non-anima max cap, not Anima's 0.5 cap diff --git a/tests/test_host_state.py b/tests/test_host_state.py index 36eed2f..b0ec2b2 100644 --- a/tests/test_host_state.py +++ b/tests/test_host_state.py @@ -1,6 +1,6 @@ import types -from ranboorux.host_state import HostMutationScope, RunContext +from ranboorux.mutation_scope import HostMutationScope, RunContext def test_host_mutation_scope_restores_changed_and_missing_attrs(): diff --git a/tests/test_lifecycle_contract.py b/tests/test_lifecycle_contract.py index 1dbaa2a..48fa37f 100644 --- a/tests/test_lifecycle_contract.py +++ b/tests/test_lifecycle_contract.py @@ -66,6 +66,8 @@ def _args(**overrides): "catalog_path": "", "lora_auto_detect_pony": True, "lora_detected_loras": [], + "anima_auto_detect": False, + "anima_tune_img2img": True, "lora_blacklist": [], } defaults.update(overrides) @@ -162,6 +164,7 @@ def test_argument_parse_failure_releases_processing_guards(stub_modules): def test_booru_error_redacts_credential_url(stub_modules): import scripts.ranbooru as ranbooru + from ranboorux.boorus import Booru secret_url = "https://site.test/api?api_key=secret&user_id=123&tags=1girl" @@ -169,7 +172,7 @@ class FakeHttp: def get_json(self, *_args, **_kwargs): raise RuntimeError(f"boom while fetching {secret_url}") - booru = ranbooru.Booru("Gelbooru", "https://site.test") + booru = Booru("Gelbooru", "https://site.test") booru.http = FakeHttp() try: diff --git a/tests/test_modules.py b/tests/test_modules.py index b17e152..439d72d 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -1,7 +1,7 @@ def test_prompting_import(): - import ranboorux.prompting as prompting + import ranboorux.tag_pipeline as tag_pipeline - assert prompting.remove_repeated_tags("a, b, a") == "a,b" + assert tag_pipeline.remove_repeated_tags("a, b, a") == "a,b" def test_image_ops_import(): diff --git a/tests/test_prompt_and_parsing.py b/tests/test_prompt_and_parsing.py index 508efb0..b2075b1 100644 --- a/tests/test_prompt_and_parsing.py +++ b/tests/test_prompt_and_parsing.py @@ -1,18 +1,18 @@ def test_remove_repeated_tags(): - import scripts.ranbooru as ranbooru + from ranboorux import tag_pipeline - assert ranbooru.remove_repeated_tags("a, b, a, c") == "a,b,c" - assert ranbooru.remove_repeated_tags("") == "" - assert ranbooru.remove_repeated_tags(None) == "" + assert tag_pipeline.remove_repeated_tags("a, b, a, c") == "a,b,c" + assert tag_pipeline.remove_repeated_tags("") == "" + assert tag_pipeline.remove_repeated_tags(None) == "" def test_limit_prompt_tags(): - import scripts.ranbooru as ranbooru + from ranboorux import tag_pipeline - assert ranbooru.limit_prompt_tags("a, b, c, d", 0.5, "Limit") == "a,b" - assert ranbooru.limit_prompt_tags("a, b, c, d", 2, "Max") == "a,b" - assert ranbooru.limit_prompt_tags("a, b", "bad", "Max") == "a, b" - assert ranbooru.limit_prompt_tags("a, b", 1, "Unknown") == "a, b" + assert tag_pipeline.limit_prompt_tags("a, b, c, d", 0.5, "Limit") == "a,b" + assert tag_pipeline.limit_prompt_tags("a, b, c, d", 2, "Max") == "a,b" + assert tag_pipeline.limit_prompt_tags("a, b", "bad", "Max") == "a, b" + assert tag_pipeline.limit_prompt_tags("a, b", 1, "Unknown") == "a, b" def test_sanitize_gelbooru_credential_variants(): @@ -35,9 +35,9 @@ def test_sanitize_gelbooru_compat_base_url(): def test_gelbooru_compat_parse_json_entities(): - import scripts.ranbooru as ranbooru + from ranboorux.boorus.gelbooru import GelbooruCompatible - client = ranbooru.GelbooruCompatible("https://example.com") + client = GelbooruCompatible("https://example.com") payload = {"post": [{"id": "1"}, {"id": "2"}], "@attributes": {"count": "42"}} entries, approx = client._parse_json_entities(payload, "post") assert [entry["id"] for entry in entries] == ["1", "2"] @@ -45,9 +45,9 @@ def test_gelbooru_compat_parse_json_entities(): def test_gelbooru_compat_parse_xml_entities(): - import scripts.ranbooru as ranbooru + from ranboorux.boorus.gelbooru import GelbooruCompatible - client = ranbooru.GelbooruCompatible("https://example.com") + client = GelbooruCompatible("https://example.com") xml_payload = ( "" "" @@ -60,9 +60,9 @@ def test_gelbooru_compat_parse_xml_entities(): def test_standardize_post_uses_tag_dict(): - import scripts.ranbooru as ranbooru + from ranboorux.boorus import Booru - booru = ranbooru.Booru("Test", "https://example.com") + booru = Booru("Test", "https://example.com") post = booru._standardize_post( { "tags": {"artist": ["alice"], "character": ["bob"], "copyright": ["copy"]}, @@ -78,9 +78,9 @@ def test_standardize_post_uses_tag_dict(): def test_standardize_post_tag_string_override_and_heuristic(): - import scripts.ranbooru as ranbooru + from ranboorux.boorus import Booru - booru = ranbooru.Booru("Test", "https://example.com") + booru = Booru("Test", "https://example.com") post = booru._standardize_post( { "tags": "foo_(series) bar", @@ -114,7 +114,7 @@ def test_show_fringe_benefits_only_visible_for_gelbooru(active_gradio_version): assert danbooru is not None -def test_loranado_scan_detects_ponyxl_and_anima_markers(tmp_path): +def test_loranado_scan_detects_ponyxl_markers(tmp_path): import types import scripts.ranbooru as ranbooru @@ -123,12 +123,6 @@ def test_loranado_scan_detects_ponyxl_and_anima_markers(tmp_path): _write_dummy_safetensors(tmp_path / "pony_magic.safetensors") _write_dummy_safetensors(tmp_path / "xlp_style.safetensors") _write_dummy_safetensors(tmp_path / "ponytail_style.safetensors") - _write_dummy_safetensors(tmp_path / "anima_style.safetensors") - _write_dummy_safetensors(tmp_path / "animapencil_v1.safetensors") - _write_dummy_safetensors(tmp_path / "anima-xl_character.safetensors") - _write_dummy_safetensors(tmp_path / "animal_ears.safetensors") - _write_dummy_safetensors(tmp_path / "animation_test.safetensors") - _write_dummy_safetensors(tmp_path / "anime_girl.safetensors") _write_dummy_safetensors( tmp_path / "metadata_style.safetensors", metadata={"ss_base_model_version": "PonyDiffusionXL"}, @@ -137,13 +131,9 @@ def test_loranado_scan_detects_ponyxl_and_anima_markers(tmp_path): tmp_path / "metadata_arch.safetensors", metadata={"modelspec.architecture": "Pony XL"}, ) - _write_dummy_safetensors( - tmp_path / "metadata_anima.safetensors", - metadata={"modelspec.architecture": "Anima Pencil XL"}, - ) _write_dummy_safetensors( tmp_path / "metadata_noise.safetensors", - metadata={"ss_tag_frequency": {"pony": 3, "anima": 5}}, + metadata={"ss_tag_frequency": {"pony": 3}}, ) _write_dummy_safetensors( tmp_path / "generic_style.safetensors", @@ -154,17 +144,10 @@ def test_loranado_scan_detects_ponyxl_and_anima_markers(tmp_path): result = script._scan_loranado_candidates("") assert "pony_magic" in result["detected_names"] assert "xlp_style" in result["detected_names"] - assert "anima_style" in result["detected_names"] - assert "animapencil_v1" in result["detected_names"] - assert "anima-xl_character" in result["detected_names"] assert "metadata_style" in result["detected_names"] assert "metadata_arch" in result["detected_names"] - assert "metadata_anima" in result["detected_names"] assert "generic_style" not in result["detected_names"] assert "ponytail_style" not in result["detected_names"] - assert "animal_ears" not in result["detected_names"] - assert "animation_test" not in result["detected_names"] - assert "anime_girl" not in result["detected_names"] assert "metadata_noise" not in result["detected_names"] @@ -180,10 +163,6 @@ def test_loranado_detection_ignores_unrelated_metadata_keys(): "generic_style.safetensors", {"ss_base_model_version": "Pony XL"}, ) - assert script._is_ponyxl_lora( - "generic_style.safetensors", - {"modelspec.architecture": "Anima Pencil XL"}, - ) def test_apply_loranado_respects_enabled_and_blacklist(tmp_path): diff --git a/tests/test_requesting.py b/tests/test_requesting.py index 662fa53..93227b5 100644 --- a/tests/test_requesting.py +++ b/tests/test_requesting.py @@ -1,12 +1,12 @@ import json import types -from ranboorux import requesting +from ranboorux import http_client def _public_dns(monkeypatch): monkeypatch.setattr( - requesting.socket, + http_client.socket, "getaddrinfo", lambda *_args, **_kwargs: [ (None, None, None, None, ("93.184.216.34", 443)), @@ -17,7 +17,7 @@ def _public_dns(monkeypatch): def test_redact_url_hides_credential_query_values(): url = "https://site.test/api?api_key=secret&user_id=123&tags=1girl" - assert requesting.redact_url(url) == ( + assert http_client.redact_url(url) == ( "https://site.test/api?api_key=&user_id=&tags=1girl" ) @@ -45,9 +45,9 @@ def get(self, *_args, **_kwargs): AssertionError("global install_cache should not be called") ), ) - monkeypatch.setattr(requesting, "requests_cache", fake_cache) + monkeypatch.setattr(http_client, "requests_cache", fake_cache) - session = requesting.BooruSession(use_cache=True) + session = http_client.BooruSession(use_cache=True) assert isinstance(session._session, FakeCachedSession) assert calls @@ -66,12 +66,12 @@ def get(self, *_args, **_kwargs): close=lambda: None, ) - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get_bytes("https://site.test/image.png", max_bytes=4) - except requesting.ResponseTooLargeError as exc: + except http_client.ResponseTooLargeError as exc: assert "exceeded 4 bytes" in str(exc) else: raise AssertionError("expected ResponseTooLargeError") @@ -85,12 +85,12 @@ def get(self, *_args, **_kwargs): calls.append(True) return types.SimpleNamespace(status_code=200, headers={}) - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get("http://127.0.0.1/private") - except requesting.UnsafeUrlError: + except http_client.UnsafeUrlError: pass else: raise AssertionError("expected UnsafeUrlError") @@ -100,8 +100,8 @@ def get(self, *_args, **_kwargs): def test_validate_outbound_url_rejects_carrier_grade_nat(): try: - requesting.validate_outbound_url("http://100.64.0.1/api") - except requesting.UnsafeUrlError: + http_client.validate_outbound_url("http://100.64.0.1/api") + except http_client.UnsafeUrlError: pass else: raise AssertionError("expected UnsafeUrlError") @@ -109,8 +109,8 @@ def test_validate_outbound_url_rejects_carrier_grade_nat(): def test_validate_outbound_url_rejects_ipv6_loopback(): try: - requesting.validate_outbound_url("http://[::1]/api") - except requesting.UnsafeUrlError: + http_client.validate_outbound_url("http://[::1]/api") + except http_client.UnsafeUrlError: pass else: raise AssertionError("expected UnsafeUrlError") @@ -118,7 +118,7 @@ def test_validate_outbound_url_rejects_ipv6_loopback(): def test_validate_outbound_url_rejects_hostname_alias_with_private_result(monkeypatch): monkeypatch.setattr( - requesting.socket, + http_client.socket, "getaddrinfo", lambda *_args, **_kwargs: [ (None, None, None, None, ("93.184.216.34", 443)), @@ -127,8 +127,8 @@ def test_validate_outbound_url_rejects_hostname_alias_with_private_result(monkey ) try: - requesting.validate_outbound_url("https://site.test/api") - except requesting.UnsafeUrlError: + http_client.validate_outbound_url("https://site.test/api") + except http_client.UnsafeUrlError: pass else: raise AssertionError("expected UnsafeUrlError") @@ -148,8 +148,8 @@ def close(self): sock = FakeSocket() try: - requesting._validate_connected_socket(sock) - except requesting.UnsafeUrlError: + http_client._validate_connected_socket(sock) + except http_client.UnsafeUrlError: pass else: raise AssertionError("expected UnsafeUrlError") @@ -173,12 +173,12 @@ def get(self, url, **_kwargs): ) fake = FakeSession() - monkeypatch.setattr(requesting.requests, "Session", lambda: fake) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: fake) + session = http_client.BooruSession(use_cache=False) try: session.get("https://site.test/start") - except requesting.UnsafeUrlError: + except http_client.UnsafeUrlError: pass else: raise AssertionError("expected UnsafeUrlError") @@ -208,14 +208,14 @@ def get(self, url, **_kwargs): uncached_calls.append(url) return FakeResponse() - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeUncachedSession()) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeUncachedSession()) monkeypatch.setattr( - requesting, + http_client, "requests_cache", types.SimpleNamespace(CachedSession=FakeCachedSession), ) - session = requesting.BooruSession(use_cache=True) + session = http_client.BooruSession(use_cache=True) session.get("https://site.test/api?api_key=secret&user_id=123") assert cached_calls == [] @@ -244,14 +244,14 @@ def get(self, url, **_kwargs): uncached_calls.append(url) return types.SimpleNamespace(status_code=200, headers={}) - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeUncachedSession()) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeUncachedSession()) monkeypatch.setattr( - requesting, + http_client, "requests_cache", types.SimpleNamespace(CachedSession=FakeCachedSession), ) - session = requesting.BooruSession(use_cache=True) + session = http_client.BooruSession(use_cache=True) session.get("https://site.test/start") assert cached_calls == ["https://site.test/start"] @@ -286,12 +286,12 @@ class FakeSession: def get(self, *_args, **_kwargs): return response - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get_bytes("https://site.test/image.png", max_bytes=5) - except requesting.ResponseTooLargeError: + except http_client.ResponseTooLargeError: pass else: raise AssertionError("expected ResponseTooLargeError") @@ -311,12 +311,12 @@ def get(self, *_args, **_kwargs): close=lambda: None, ) - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get_bytes("https://site.test/not-image") - except requesting.InvalidContentTypeError as exc: + except http_client.InvalidContentTypeError as exc: assert "text/html" in str(exc) else: raise AssertionError("expected InvalidContentTypeError") @@ -344,8 +344,8 @@ class FakeSession: def get(self, *_args, **_kwargs): return FakeResponse() - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) assert session.get_json("https://site.test/api") == {"ok": True} @@ -370,12 +370,12 @@ class FakeSession: def get(self, *_args, **_kwargs): return FakeResponse() - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get_json("https://site.test/api", max_bytes=10) - except requesting.ResponseTooLargeError: + except http_client.ResponseTooLargeError: pass else: raise AssertionError("expected ResponseTooLargeError") @@ -407,12 +407,12 @@ class FakeSession: def get(self, *_args, **_kwargs): return response - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get_json("https://site.test/api", max_bytes=6) - except requesting.ResponseTooLargeError: + except http_client.ResponseTooLargeError: pass else: raise AssertionError("expected ResponseTooLargeError") @@ -437,12 +437,12 @@ class FakeSession: def get(self, *_args, **_kwargs): return FakeResponse() - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get_json("https://site.test/api") - except requesting.InvalidContentTypeError: + except http_client.InvalidContentTypeError: pass else: raise AssertionError("expected InvalidContentTypeError") @@ -469,8 +469,8 @@ class FakeSession: def get(self, *_args, **_kwargs): return FakeResponse() - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) try: session.get_json("https://site.test/api") @@ -482,7 +482,7 @@ def get(self, *_args, **_kwargs): def test_redact_url_mixed_case_and_percent_encoding(): url1 = "https://site.test/api?X-Amz-Signature=secret123&X-goog-Credential=secret456&sig=secret789&normal=hello" - redacted1 = requesting.redact_url(url1) + redacted1 = http_client.redact_url(url1) assert "secret123" not in redacted1 assert "secret456" not in redacted1 assert "secret789" not in redacted1 @@ -492,7 +492,7 @@ def test_redact_url_mixed_case_and_percent_encoding(): assert "normal=hello" in redacted1 url2 = "https://site.test/api?api_key=secret%20key&password=hello%26world" - redacted2 = requesting.redact_url(url2) + redacted2 = http_client.redact_url(url2) assert "secret%20key" not in redacted2 assert "hello%26world" not in redacted2 assert "api_key=" in redacted2 @@ -504,7 +504,7 @@ def test_exception_sanitization_mixed_content(): "Error accessing file E:\\private\\forge\\extensions\\sd_forge_controlnet " "when calling https://cdn.test/foo?X-Amz-Signature=supersecret&normal=param" ) - sanitized = requesting.sanitize_exception_text(exc_text) + sanitized = http_client.sanitize_exception_text(exc_text) assert "E:\\private" not in sanitized assert "supersecret" not in sanitized assert "" in sanitized @@ -569,14 +569,14 @@ class FakeUncachedSession: def get(self, url, **kwargs): return FakeResponse(200, {}) - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeUncachedSession()) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeUncachedSession()) monkeypatch.setattr( - requesting, + http_client, "requests_cache", types.SimpleNamespace(CachedSession=FakeCachedSession), ) - session = requesting.BooruSession(use_cache=True) + session = http_client.BooruSession(use_cache=True) session.get("https://site.test/start") assert not session._session.cache.contains("https://site.test/start") @@ -600,7 +600,7 @@ def test_sanitize_exception_invalid_url_with_secrets_and_paths(): exc = InvalidURL(err_msg) - sanitized_exc = requesting.sanitize_exception(exc) + sanitized_exc = http_client.sanitize_exception(exc) assert isinstance(sanitized_exc, RuntimeError) message = str(sanitized_exc) @@ -638,8 +638,8 @@ class FakeSession: def get(self, *_args, **_kwargs): return FakeResponse() - monkeypatch.setattr(requesting.requests, "Session", lambda: FakeSession()) - session = requesting.BooruSession(use_cache=False) + monkeypatch.setattr(http_client.requests, "Session", lambda: FakeSession()) + session = http_client.BooruSession(use_cache=False) def mock_loads(*_args, **_kwargs): raise json.JSONDecodeError( @@ -648,11 +648,11 @@ def mock_loads(*_args, **_kwargs): 0, ) - monkeypatch.setattr(requesting.json, "loads", mock_loads) + monkeypatch.setattr(http_client.json, "loads", mock_loads) try: session.get_json("https://site.test/api") - except requesting.BooruResponseError as exc: + except http_client.BooruResponseError as exc: message = str(exc) assert exc.__cause__ is not None assert isinstance(exc.__cause__, json.JSONDecodeError) @@ -671,16 +671,16 @@ def test_path_redaction_with_spaces(): posix_path = "/home/user/Private Folder/file.py" file_path = "file:" + "///C:/Users/user/Private Folder/file.py" - assert "user profile" not in requesting.sanitize_exception_text( + assert "user profile" not in http_client.sanitize_exception_text( "Error C:\\Users\\user profile\\Private Folder\\file.py." ) - assert "share name" not in requesting.sanitize_exception_text( + assert "share name" not in http_client.sanitize_exception_text( "Error \\\\server\\share name\\folder\\file.py." ) - assert "Private Folder" not in requesting.sanitize_exception_text( + assert "Private Folder" not in http_client.sanitize_exception_text( "Error /home/user/Private Folder/file.py." ) - assert "Private Folder" not in requesting.sanitize_exception_text( + assert "Private Folder" not in http_client.sanitize_exception_text( "Error file:" + "///C:/Users/user/Private Folder/file.py." ) @@ -690,7 +690,7 @@ def test_path_redaction_with_spaces(): f"and {file_path} with signed URL {signed_url}" ) - sanitized = requesting.sanitize_exception_text(mixed_msg) + sanitized = http_client.sanitize_exception_text(mixed_msg) assert "user profile" not in sanitized assert "share name" not in sanitized diff --git a/tests/test_run_options.py b/tests/test_run_options.py index 5be8c5b..4a8dde1 100644 --- a/tests/test_run_options.py +++ b/tests/test_run_options.py @@ -4,7 +4,7 @@ def test_ui_argument_field_order_is_frozen(): - assert len(UI_ARGUMENT_FIELDS) == 62 + assert len(UI_ARGUMENT_FIELDS) == 64 assert UI_ARGUMENT_FIELDS[:6] == ( "enabled", "tags", @@ -13,12 +13,13 @@ def test_ui_argument_field_order_is_frozen(): "gelbooru_user_id", "gelbooru_compat_base_url", ) - assert UI_ARGUMENT_FIELDS[-5:] == ( - "use_tag_catalog", + assert UI_ARGUMENT_FIELDS[-6:] == ( "catalog_path", "lora_auto_detect_pony", "lora_detected_loras", "lora_blacklist", + "anima_auto_detect", + "anima_tune_img2img", ) @@ -33,11 +34,12 @@ def test_run_options_from_script_args_maps_names_once(): assert options.image_workflow.use_img2img == 12 assert options.tag_filters.remove_text_tags == 50 assert options.loranado.blacklist == 61 + assert options.anima_tune_img2img == 63 assert options.as_dict() == dict(zip(UI_ARGUMENT_FIELDS, values)) def test_run_options_rejects_wrong_count(): - with pytest.raises(ValueError, match="Expected 62"): + with pytest.raises(ValueError, match="Expected 64"): RunOptions.from_script_args([object()]) diff --git a/tests/test_tag_pipeline.py b/tests/test_tag_pipeline.py index a84a40c..78f3603 100644 --- a/tests/test_tag_pipeline.py +++ b/tests/test_tag_pipeline.py @@ -215,3 +215,108 @@ def test_post_rejected_by_filter(): favorites_guard=set(), ) assert rejected is False + + +def test_post_rejected_by_filter_remove_furry(): + """Test that remove_furry flag rejects furry tags.""" + post = {"id": "1", "booru_name": "danbooru", "tags": "kemonomimi, 1girl, blonde_hair"} + cache = {} + rejected, reason = post_rejected_by_filter( + post, + filter_ctx=None, + toggles=(False, False, False, False, False, True, False, False, False, False), + base_colors=(set(), set()), + allowed_subjects=set(), + cache=cache, + favorites_guard=set(), + ) + assert rejected is True + assert reason["rule"] == "furry" + + +def test_post_rejected_by_filter_remove_clothing(): + """Test that remove_clothing rejects clothing tags but not 'no_clothing'.""" + post = {"id": "2", "booru_name": "danbooru", "tags": "dress, 1girl, no_clothing"} + cache = {} + rejected, reason = post_rejected_by_filter( + post, + filter_ctx=None, + toggles=(False, False, True, False, False, False, False, False, False, False), + base_colors=(set(), set()), + allowed_subjects=set(), + cache=cache, + favorites_guard=set(), + ) + assert rejected is True + assert reason["rule"] == "clothing" + + +def test_post_rejected_by_filter_remove_headwear(): + """Test remove_headwear with halo edge case.""" + post = {"id": "3", "booru_name": "danbooru", "tags": "halo, 1girl, blonde_hair"} + cache = {} + rejected, reason = post_rejected_by_filter( + post, + filter_ctx=None, + toggles=(False, False, False, False, False, False, True, False, False, False), + base_colors=(set(), set()), + allowed_subjects=set(), + cache=cache, + favorites_guard=set(), + ) + assert rejected is True + assert reason["rule"] == "headwear" + + +def test_post_rejected_by_filter_remove_girl_suffix(): + """Test remove_girl_suffix rejects _girl tags but not 1girl/girl.""" + post = {"id": "4", "booru_name": "danbooru", "tags": "cat_girl, 1girl, girl, blonde_hair"} + cache = {} + rejected, reason = post_rejected_by_filter( + post, + filter_ctx=None, + toggles=(False, False, False, False, False, False, False, True, False, False), + base_colors=(set(), set()), + allowed_subjects=set(), + cache=cache, + favorites_guard=set(), + ) + assert rejected is True + assert reason["rule"] == "girl-suffix" + assert reason["tag"] == "cat_girl" + + +def test_post_rejected_by_filter_remove_character(): + """Test remove_character rejects character tags.""" + post = {"id": "5", "booru_name": "danbooru", "tags": "1girl", "character_tags": "heroine"} + cache = {} + rejected, reason = post_rejected_by_filter( + post, + filter_ctx=None, + toggles=(False, True, False, False, False, False, False, False, False, False), + base_colors=(set(), set()), + allowed_subjects=set(), + cache=cache, + favorites_guard=set(), + ) + assert rejected is True + assert reason["rule"] == "character" + + +def test_post_rejected_by_filter_favorites_guard(): + """Test that favorites_guard bypasses removal matching.""" + post = {"id": "6", "booru_name": "danbooru", "tags": "bad_tag, 1girl"} + removal_raw = ["bad_tag"] + ctx = build_removal_context(removal_raw, favorites_raw=[], synonym_lookup={}) + cache = {} + # With favorites_guard containing "bad_tag" - should NOT be rejected + rejected, reason = post_rejected_by_filter( + post, + filter_ctx=ctx, + toggles=(False, False, False, False, False, False, False, False, False, False), + base_colors=(set(), set()), + allowed_subjects=set(), + cache=cache, + favorites_guard={"bad tag"}, + ) + assert rejected is False diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py index c0c294e..849402b 100644 --- a/tests/test_wrappers.py +++ b/tests/test_wrappers.py @@ -1,12 +1,14 @@ def test_prompt_wrappers_match_module(): import scripts.ranbooru as ranbooru - from ranboorux import prompting + from ranboorux import tag_pipeline prompt = "a, b, a, c" - assert ranbooru.remove_repeated_tags(prompt) == prompting.remove_repeated_tags(prompt) - assert ranbooru.limit_prompt_tags("a, b, c, d", 2, "Max") == prompting.limit_prompt_tags( + assert ranbooru.rb_tag_pipeline.remove_repeated_tags( + prompt + ) == tag_pipeline.remove_repeated_tags(prompt) + assert ranbooru.rb_tag_pipeline.limit_prompt_tags( "a, b, c, d", 2, "Max" - ) + ) == tag_pipeline.limit_prompt_tags("a, b, c, d", 2, "Max") def test_controlnet_wrapper_uses_integration(monkeypatch): diff --git a/tools/build_release.py b/tools/build_release.py index 1100893..6f8e7b9 100644 --- a/tools/build_release.py +++ b/tools/build_release.py @@ -15,7 +15,7 @@ "pyproject.toml", "README.md", "requirements.txt", - "adetailer/**/*", + # "adetailer/**/*", # local nested extension dir (ignored by .gitignore); do not package "data/**/*", "docs/CHANGELOG.md", "docs/CONFIG.md", @@ -68,8 +68,8 @@ ".txt", ".yml", } -LOCAL_FILE_URI_RE = re.compile(rb"file:" + rb"///", re.IGNORECASE) -WINDOWS_ABSOLUTE_PATH_RE = re.compile(rb"\b[A-Za-z]:\\[^\\\r\n\t ]+\\[^\\\r\n\t ]+") +LOCAL_FILE_URI_RE = re.compile(rb"\bfile:///") +WINDOWS_ABSOLUTE_PATH_RE = re.compile(rb"\b[A-Za-z]:(?:\\[^\\\r\n\t]+)+") def matches_any(path, patterns): @@ -192,7 +192,6 @@ def run_self_tests(): "README.md", "install.py", "scripts/ranbooru.py", - "ranboorux/prompting.py", "docs/usage.md", ] for f in allowed: diff --git a/tools/inspect_ui.py b/tools/inspect_ui.py index 9289d24..2c82de4 100644 --- a/tools/inspect_ui.py +++ b/tools/inspect_ui.py @@ -158,7 +158,7 @@ def __exit__(self, exc_type, exc, tb): os.makedirs(os.path.dirname(output_path), exist_ok=True) # We map components back to their indices and variable names -# The return statement from scripts/ranbooru.py has 62 items: +# The return statement from scripts/ranbooru.py has 64 items: variable_names = [ "enabled", "tags", @@ -216,12 +216,13 @@ def __exit__(self, exc_type, exc, tb): "remove_headwear_tags", "remove_girl_suffix_tags", "preserve_hair_eye_colors", - "remove_series_tags", "use_tag_catalog", "catalog_path", "lora_auto_detect_pony", "lora_detected_loras", "lora_blacklist", + "anima_auto_detect", + "anima_tune_img2img", ] with open(output_path, "w", encoding="utf-8") as f: diff --git a/tools/repo_guard.py b/tools/repo_guard.py index d6b8205..75632d6 100644 --- a/tools/repo_guard.py +++ b/tools/repo_guard.py @@ -33,13 +33,18 @@ def check_files(file_list: Iterable[str], forbidden_prefixes: Sequence[str]) -> def _run_git(repo_root: Path, args: Sequence[str]) -> subprocess.CompletedProcess[str]: - return subprocess.run( - ["git", "-c", f"safe.directory={repo_root.as_posix()}", *args], - cwd=repo_root, - capture_output=True, - text=True, - check=True, - ) + try: + return subprocess.run( + ["git", "-c", "safe.directory=*", *args], + cwd=repo_root, + capture_output=True, + text=True, + check=True, + ) + except (subprocess.CalledProcessError, FileNotFoundError): + return subprocess.CompletedProcess( + args=["git", *args], returncode=128, stdout="", stderr="" + ) def get_git_modified_files(repo_root: Path) -> List[str]: