KBlueLeaf/DanTagGen-gamma

DanTagGen - gamma

DanTagGen(Danbooru Tag Generator) is inspired from p1atdev's dart project. But with different arch, dataset, format and different training strategy.

Difference between versions

alpha: pretrain on 2M dataset, smaller batch size. Limited ability
beta: pretrain on 5.3M dataset, larger batch size. More stable, better ability with only a few information provided.
gamma: finetuned from beta, with 3.6M dataset (union of all posts after id 5,000,000 and top25% fav count posts)

Model arch

This version of DTG is trained from scratch with 400M param LLaMA arch.(In my personal preference I will call it NanoLLaMA) Since it is llama arch. Theoritically it should be able to be used in any LLaMA inference interface.

This repo also provided converted FP16 gguf model and quantized 8bit/6bit gguf models. Basically it is recommended to use llama.cpp or llama-cpp-python to run this model. Which will be very fast.

Format

prompt = f"""
rating: {rating or '<|empty|>'}
artist: {artist.strip() or '<|empty|>'}
characters: {characters.strip() or '<|empty|>'}
copyrights: {copyrights.strip() or '<|empty|>'}
aspect ratio: {f"{aspect_ratio:.1f}" or '<|empty|>'}
target: {'<|' + target + '|>' if target else '<|long|>'}
general: {", ".join(special_tags)}, {general.strip().strip(",")}<|input_end|>
"""

for example:

rating: safe
artist: <|empty|>
characters: <|empty|>
copyrights: <|empty|>
aspect ratio: 1.0
target: <|short|>
general: 1girl, solo, dragon girl, dragon horns, dragon tail<|input_end|>

And you may get something like:

rating: safe
artist: <|empty|>
characters: <|empty|>
copyrights: <|empty|>
aspect ratio: 1.0
target: <|short|>
general: 1girl, solo, dragon girl, dragon horns, dragon tail<|input_end|>open mouth, red eyes, long hair, pointy ears, tail, black hair, chinese clothes, simple background, dragon, hair between eyes, horns, china dress, dress, looking at viewer, breasts

Utilities

HF space: https://huggingface.co/spaces/KBlueLeaf/DTG-demo
SD-WebUI extension (Forge compatible): https://github.com/KohakuBlueleaf/z-a1111-sd-webui-dtg
Third Party ComfyUI Node: https://github.com/toyxyz/a1111-sd-webui-dtg_comfyui