A simple FastAPI Server to run WD1.4-Tagger

Feel free to make PRs or use the code for your own needs

Changelog

You can keep track of all changes on the release page

Installation

Simple installation :

pip install wd14-tagger-api

This will install all the necessary dependencies, including a CPU support only version of PyTorch

I recommend that you install the GPU version to improve processing speed.

Don't forget to add the -d gpu flag

Note The first startup using a video card may take a long time ( up to 30-60 seconds ) but then the result of processing will be almost minimal

Windows

python -m venv venv
venv\Scripts\activate
pip install wd14-tagger-api
pip install onnxruntime-gpu

Linux

python -m venv venv
source venv\bin\activate
pip install wd14-tagger-api
pip install onnxruntime-gpu

Manual

# Clone REPO
git clone https://github.com/daswer123/wd14-tagger-api-server
cd wd14-tagger-api-server
# Create virtual env
python -m venv venv
venv/scripts/activate or source venv/bin/activate
# Install deps
pip install -r requirements.txt
pip install onnxruntime-gpu
# Launch server
python -m wd14_tagger_api

Starting Server

python -m wd14_tagger_api will run on default ip and port (localhost:8019)

usage: xtts_api_server [-h] [-hs HOST] [-p PORT] [-d DEVICE] [-wdm WD14 model] [-wdt WD14 threshold] [--replace-underscore]

Run WD14-Tagger within a FastAPI application

options:
  -h, --help show this help message and exit
  -hs HOST, --host HOST
  -p PORT, --port PORT
  -d DEVICE, --device DEVICE `cpu` or `gpu`
  -wdm WD14 Model, --wd14-model You can specify which model to use when loading the server.
  -wdt WD14 Threshold, --wd14-threshold  Threshold value for WD14
  --replace-underscore Will return spaces instead of underscore

Avalible Models

wd14-vit.v2
wd14-vit.v2
wd14-convnext.v1
wd14-convnext.v2
wd14-convnextv2.v1
wd14-swinv2-v1
wd-v1-4-moat-tagger.v2
mld-caformer.dec-5-97527
mld-tresnetd.6-30000

API Docs

API documentation can be accessed from http://localhost:8002/docs when the server is running. This documentation is automatically generated by FastAPI and provides an interactive interface to test the API endpoints.

Endpoints Overview

The server provides several endpoints related to image tagging functionality:

POST `/update-settings/`

This endpoint allows updating model settings dynamically.

Request Body: JSON object containing new settings values.
- wd14_model: String (optional). The name of the model to use.
- threshold: Float (optional). The threshold value for tag generation.

Example Request:

{
  "wd14_model": "new-model-name",
  "threshold": 0.5
}

Response: A confirmation message indicating that settings have been updated successfully.

POST `/tag-image/`

Upload an image file to get associated tags based on current model and threshold settings.

Form Data:
- file: File. Image file to be tagged.

Response: A string of comma-separated tags generated for the uploaded image. If WD14_REPLACE_UNDERSCORE environment variable is set to true, underscores in tags are replaced with spaces.

Usage Examples

Below are examples using curl command-line tool to interact with the server:

Updating Model Settings

curl -X 'POST' \
'http://localhost:8002/update-settings/' \
-H 'accept: application/json' \
-H 'Content-Type: application/json' \
-d '{"wd14_model": "new-model-name", "threshold": 0.5}'

Tagging an Image

curl -X 'POST' \
'http://localhost:8002/tag-image/' \
-H 'accept: application/json' \
-F 'file=@path_to_your_image.jpg;type=image/jpeg'

Replace "path_to_your_image.jpg" with the actual path to your image file you wish to tag.

Credit

Thanks corkborg for its standalone version, I used it as a base for the server.

daswer123/wd14-tagger-api-server