Upload 6 files

wdv3-timm-main/.gitignore

@@ -0,0 +1,253 @@
+# Created by https://www.toptal.com/developers/gitignore/api/linux,windows,macos,visualstudiocode,python
+# Edit at https://www.toptal.com/developers/gitignore?templates=linux,windows,macos,visualstudiocode,python
+
+### Linux ###
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+### macOS ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+### VisualStudioCode ###
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+*.code-workspace
+
+# Local History for Visual Studio Code
+.history/
+
+### VisualStudioCode Patch ###
+# Ignore all local history of files
+.history
+.ionide
+
+### Windows ###
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+# End of https://www.toptal.com/developers/gitignore/api/linux,windows,macos,visualstudiocode,python
+
+# temp and misc
+/misc/
+/temp/
+
+# direnv
+.envrc
+.envrc.*
+
+# dotenv
+.env
+.env.*
+
+# temp files
+**/tmp_*.*
+**/*.tmp.*
+
+# but keep examples
+!*.example
+
+# input images and heatmap outputs
+/images/
+/heatmaps/

wdv3-timm-main/.vscode/settings.json

@@ -0,0 +1,94 @@
+{
+    "editor.insertSpaces": true,
+    "editor.tabSize": 4,
+    "files.trimTrailingWhitespace": true,
+    "editor.rulers": [100, 120],
+
+    "files.associations": {
+        "*.yaml": "yaml"
+    },
+    "files.exclude": {
+        "**/.git": true,
+        "**/.svn": true,
+        "**/.hg": true,
+        "**/CVS": true,
+        "**/.DS_Store": true,
+        "**/Thumbs.db": true,
+        "**/.ruff_cache": true,
+        "**/__pycache__": true,
+        "**/*.egg-info": true
+    },
+
+    "[shellscript]": {
+        "files.eol": "\n",
+        "editor.tabSize": 4,
+        "editor.detectIndentation": false
+    },
+
+    "[python]": {
+        "editor.wordBasedSuggestions": "off",
+        "editor.formatOnSave": true,
+        "editor.defaultFormatter": "charliermarsh.ruff",
+        "editor.codeActionsOnSave": {
+            "source.organizeImports": "always"
+        }
+    },
+    "python.analysis.include": ["./src", "./scripts", "./tests"],
+
+    "[json]": {
+        "editor.defaultFormatter": "esbenp.prettier-vscode",
+        "editor.detectIndentation": false,
+        "editor.formatOnSaveMode": "file",
+        "editor.formatOnSave": true,
+        "editor.tabSize": 2
+    },
+    "[jsonc]": {
+        "editor.defaultFormatter": "esbenp.prettier-vscode",
+        "editor.detectIndentation": false,
+        "editor.formatOnSaveMode": "file",
+        "editor.formatOnSave": true,
+        "editor.tabSize": 2
+    },
+
+    "[toml]": {
+        "editor.tabSize": 2,
+        "editor.detectIndentation": false,
+        "editor.formatOnSave": true,
+        "editor.formatOnSaveMode": "file",
+        "editor.defaultFormatter": "tamasfe.even-better-toml",
+        "editor.rulers": [80, 100]
+    },
+    "evenBetterToml.formatter.columnWidth": 88,
+
+    "[yaml]": {
+        "editor.detectIndentation": false,
+        "editor.tabSize": 2,
+        "editor.formatOnSave": true,
+        "editor.formatOnSaveMode": "file",
+        "diffEditor.ignoreTrimWhitespace": false,
+        "editor.defaultFormatter": "redhat.vscode-yaml"
+    },
+    "yaml.format.bracketSpacing": true,
+    "yaml.format.proseWrap": "preserve",
+    "yaml.format.singleQuote": false,
+    "yaml.format.printWidth": 110,
+
+    "[hcl]": {
+        "editor.detectIndentation": false,
+        "editor.formatOnSave": true,
+        "editor.formatOnSaveMode": "file",
+        "editor.defaultFormatter": "fredwangwang.vscode-hcl-format"
+    },
+
+    "[markdown]": {
+        "files.trimTrailingWhitespace": false
+    },
+
+    "css.lint.validProperties": ["dock", "content-align", "content-justify"],
+    "[css]": {
+        "editor.formatOnSave": true
+    },
+
+    "remote.autoForwardPorts": false,
+    "remote.autoForwardPortsSource": "process"
+}

wdv3-timm-main/README.md

@@ -0,0 +1,84 @@
+# wdv3-timm
+
+small example thing showing how to use `timm` to run the WD Tagger V3 models.
+
+## How To Use
+
+1. clone the repository and enter the directory:
+```sh
+git clone https://github.com/neggles/wdv3-timm.git
+cd wd3-timm
+```
+
+2. Create a virtual environment and install the Python requirements.
+
+If you're using Linux, you can use the provided script:
+```sh
+bash setup.sh
+```
+
+Or if you're on Windows (or just want to do it manually), you can do the following:
+```sh
+# Create virtual environment
+python3.10 -m venv .venv
+# Activate it
+source .venv/bin/activate
+# Upgrade pip/setuptools/wheel
+python -m pip install -U pip setuptools wheel
+# At this point, optionally you can install PyTorch manually (e.g. if you are not using an nVidia GPU)
+python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+# Install requirements
+python -m pip install -r requirements.txt
+```
+
+3. Run the example script, picking one of the 3 models to use:
+```sh
+python wdv3_timm.py <swinv2|convnext|vit> path/to/image.png
+```
+
+Example output from `python wdv3_timm.py vit a_picture_of_ganyu.png`:
+```sh
+Loading model 'vit' from 'SmilingWolf/wd-vit-tagger-v3'...
+Loading tag list...
+Creating data transform...
+Loading image and preprocessing...
+Running inference...
+Processing results...
+--------
+Caption: 1girl, horns, solo, bell, ahoge, colored_skin, blue_skin, neck_bell, looking_at_viewer, purple_eyes, upper_body, blonde_hair, long_hair, goat_horns, blue_hair, off_shoulder, sidelocks, bare_shoulders, alternate_costume, shirt, black_shirt, cowbell, ganyu_(genshin_impact)
+--------
+Tags: 1girl, horns, solo, bell, ahoge, colored skin, blue skin, neck bell, looking at viewer, purple eyes, upper body, blonde hair, long hair, goat horns, blue hair, off shoulder, sidelocks, bare shoulders, alternate costume, shirt, black shirt, cowbell, ganyu \(genshin impact\)
+--------
+Ratings:
+  general: 0.827
+  sensitive: 0.199
+  questionable: 0.001
+  explicit: 0.001
+--------
+Character tags (threshold=0.75):
+  ganyu_(genshin_impact): 0.991
+--------
+General tags (threshold=0.35):
+  1girl: 0.996
+  horns: 0.950
+  solo: 0.947
+  bell: 0.918
+  ahoge: 0.897
+  colored_skin: 0.881
+  blue_skin: 0.872
+  neck_bell: 0.854
+  looking_at_viewer: 0.817
+  purple_eyes: 0.734
+  upper_body: 0.615
+  blonde_hair: 0.609
+  long_hair: 0.607
+  goat_horns: 0.524
+  blue_hair: 0.496
+  off_shoulder: 0.472
+  sidelocks: 0.470
+  bare_shoulders: 0.464
+  alternate_costume: 0.437
+  shirt: 0.427
+  black_shirt: 0.417
+  cowbell: 0.415
+```

wdv3-timm-main/requirements.txt

@@ -0,0 +1,11 @@
+diffusers
+huggingface-hub
+numpy
+pandas
+pillow >= 9.5.0
+simple-parsing >= 0.1.5
+timm @ git+https://github.com/huggingface/pytorch-image-models@main#egg=timm
+tokenizers
+torch >= 2.0.0
+torchvision
+transformers

wdv3-timm-main/setup.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# get the folder this script is in and make sure we're in it
+script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd -P)
+cd "${script_dir}"
+
+# make venv if not exist
+if [[ ! -d .venv ]]; then
+    echo "Creating virtual environment..."
+    python3.10 -m venv .venv
+fi
+
+# activate the venv
+source .venv/bin/activate
+
+# upgrade pip
+python -m pip install -U pip setuptools wheel
+
+# install requirements
+python -m pip install -r requirements.txt
+
+echo "Setup complete. Run 'source .venv/bin/activate' to enter the virtual environment."
+exit 0

wdv3-timm-main/wdv3_timm.py

@@ -0,0 +1,203 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+import timm
+import torch
+from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import HfHubHTTPError
+from PIL import Image
+from simple_parsing import field, parse_known_args
+from timm.data import create_transform, resolve_data_config
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+import json
+
+torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL_REPO_MAP = {
+    "vit": "SmilingWolf/wd-vit-tagger-v3",
+    "swinv2": "SmilingWolf/wd-swinv2-tagger-v3",
+    "convnext": "SmilingWolf/wd-convnext-tagger-v3",
+}
+
+def pil_ensure_rgb(image: Image.Image) -> Image.Image:
+    # convert to RGB/RGBA if not already (deals with palette images etc.)
+    if image.mode not in ["RGB", "RGBA"]:
+        image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB")
+    # convert RGBA to RGB with white background
+    if image.mode == "RGBA":
+        canvas = Image.new("RGBA", image.size, (255, 255, 255))
+        canvas.alpha_composite(image)
+        image = canvas.convert("RGB")
+    return image
+
+def pil_pad_square(image: Image.Image) -> Image.Image:
+    w, h = image.size
+    # get the largest dimension so we can pad to a square
+    px = max(image.size)
+    # pad to square with white background
+    canvas = Image.new("RGB", (px, px), (255, 255, 255))
+    canvas.paste(image, ((px - w) // 2, (px - h) // 2))
+    return canvas
+
+@dataclass
+class LabelData:
+    names: list[str]
+    rating: list[np.int64]
+    general: list[np.int64]
+    character: list[np.int64]
+
+def load_labels_hf(
+    repo_id: str,
+    revision: Optional[str] = None,
+    token: Optional[str] = None,
+) -> LabelData:
+    try:
+        csv_path = hf_hub_download(
+            repo_id=repo_id, filename="selected_tags.csv", revision=revision, token=token
+        )
+        csv_path = Path(csv_path).resolve()
+    except HfHubHTTPError as e:
+        raise FileNotFoundError(f"selected_tags.csv failed to download from {repo_id}") from e
+
+    df: pd.DataFrame = pd.read_csv(csv_path, usecols=["name", "category"])
+    tag_data = LabelData(
+        names=df["name"].tolist(),
+        rating=list(np.where(df["category"] == 9)[0]),
+        general=list(np.where(df["category"] == 0)[0]),
+        character=list(np.where(df["category"] == 4)[0]),
+    )
+
+    return tag_data
+
+def get_tags(
+    probs: Tensor,
+    labels: LabelData,
+    gen_threshold: float,
+    char_threshold: float,
+):
+    # Convert indices+probs to labels
+    probs = list(zip(labels.names, probs.numpy()))
+
+    # First 4 labels are actually ratings
+    rating_labels = dict([probs[i] for i in labels.rating])
+
+    # General labels, pick any where prediction confidence > threshold
+    gen_labels = [probs[i] for i in labels.general]
+    gen_labels = dict([x for x in gen_labels if x[1] > gen_threshold])
+    gen_labels = dict(sorted(gen_labels.items(), key=lambda item: item[1], reverse=True))
+
+    # Character labels, pick any where prediction confidence > threshold
+    char_labels = [probs[i] for i in labels.character]
+    char_labels = dict([x for x in char_labels if x[1] > char_threshold])
+    char_labels = dict(sorted(char_labels.items(), key=lambda item: item[1], reverse=True))
+
+    # Combine general and character labels, sort by confidence
+    combined_names = [x for x in gen_labels]
+    combined_names.extend([x for x in char_labels])
+
+    # Convert to a string suitable for use as a training caption
+    caption = ", ".join(combined_names)
+    taglist = caption.replace("_", " ").replace("(", "\(").replace(")", "\)")
+
+    return caption, taglist, rating_labels, char_labels, gen_labels
+
+@dataclass
+class ScriptOptions:
+    image_file: Path = field(positional=True)
+    model: str = field(default="vit")
+    gen_threshold: float = field(default=0.35)
+    char_threshold: float = field(default=0.75)
+
+def main(opts: ScriptOptions):
+    repo_id = MODEL_REPO_MAP.get(opts.model)
+    image_path = Path(opts.image_file).resolve()
+    if not image_path.is_file():
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    print(f"Loading model '{opts.model}' from '{repo_id}'...")
+    model: nn.Module = timm.create_model("hf-hub:" + repo_id).eval()
+    state_dict = timm.models.load_state_dict_from_hf(repo_id)
+    model.load_state_dict(state_dict)
+
+    print("Loading tag list...")
+    labels: LabelData = load_labels_hf(repo_id=repo_id)
+
+    print("Creating data transform...")
+    transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
+
+    print("Loading image and preprocessing...")
+    # get image
+    img_input: Image.Image = Image.open(image_path)
+    # ensure image is RGB
+    img_input = pil_ensure_rgb(img_input)
+    # pad to square with white background
+    img_input = pil_pad_square(img_input)
+    # run the model's input transform to convert to tensor and rescale
+    inputs: Tensor = transform(img_input).unsqueeze(0)
+    # NCHW image RGB to BGR
+    inputs = inputs[:, [2, 1, 0]]
+
+    print("Running inference...")
+    with torch.inference_mode():
+        # move model to GPU, if available
+        if torch_device.type != "cpu":
+            model = model.to(torch_device)
+            inputs = inputs.to(torch_device)
+        # run the model
+        outputs = model.forward(inputs)
+        # apply the final activation function (timm doesn't support doing this internally)
+        outputs = F.sigmoid(outputs)
+        # move inputs, outputs, và model về CPU nếu đang ở trên GPU
+        if torch_device.type != "cpu":
+            inputs = inputs.to("cpu")
+            outputs = outputs.to("cpu")
+            model = model.to("cpu")
+
+    print("Processing results...")
+    # Đọc giá trị từ config.json
+    with open('config.json', 'r') as config_file:
+        config_data = json.load(config_file)
+
+    gen_threshold = config_data.get('general_threshold', 0.35)
+    char_threshold = config_data.get('character_threshold', 0.75)
+
+    caption, taglist, ratings, character, general = get_tags(
+        probs=outputs.squeeze(0),
+        labels=labels,
+        gen_threshold=gen_threshold,
+        char_threshold=char_threshold,
+    )
+
+    print("--------")
+    print(f"Caption: {caption}")
+    print("--------")
+    print(f"Tags: {taglist}")
+
+    print("--------")
+    print("Ratings:")
+    for k, v in ratings.items():
+        print(f"  {k}: {v:.3f}")
+
+    print("--------")
+    print(f"Character tags (threshold={char_threshold}):")
+    for k, v in character.items():
+        print(f"  {k}: {v:.3f}")
+
+    print("--------")
+    print(f"General tags (threshold={gen_threshold}):")
+    for k, v in general.items():
+        print(f"  {k}: {v:.3f}")
+
+    print("Done!")
+
+
+if __name__ == "__main__":
+    opts, _ = parse_known_args(ScriptOptions)
+    if opts.model not in MODEL_REPO_MAP:
+        print(f"Available models: {list(MODEL_REPO_MAP.keys())}")
+        raise ValueError(f"Unknown model name '{opts.model}'")
+    main(opts)