library_name: transformers
license: apache-2.0
datasets:
- isek-ai/danbooru-tags-2024
tags:
- trl
- sft
- optimum
- danbooru
inference: false
Dart (Danbooru Tags Transformer) v2
This model is a fine-tuned Dart (Danbooru Tags Transformer) v2 base model that generates danbooru tags.
Demo: 🤗 Space with ZERO
Model variants
Name | Architecture | Param size | Type |
---|---|---|---|
v2-moe-sft | Mixtral | 166m | SFT |
v2-moe-base | Mixtral | 166m | Pretrain |
v2-sft | Mistral | 114m | SFT |
v2-base | Mistral | 114m | Pretrain |
v2-vectors | Embedding | - | Tag Embedding |
Usage
Using 🤗Transformers
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL_NAME = "p1atdev/dart-v2-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16)
prompt = (
f"<|bos|>"
f"<copyright>vocaloid</copyright>"
f"<character>hatsune miku</character>"
f"<|rating:general|><|aspect_ratio:tall|><|length:long|>"
f"<general>1girl"
)
inputs = tokenizer(prompt, return_tensors="pt").input_ids
with torch.no_grad():
outputs = model.generate(
inputs,
do_sample=True,
temperature=1.0,
top_p=1.0,
top_k=100,
max_new_tokens=128,
num_beams=1,
)
print(", ".join([tag for tag in tokenizer.batch_decode(outputs[0], skip_special_tokens=True) if tag.strip() != ""]))
Using 📦dartrs
library
This library is very experimental and there will be breaking changes in the future.
📦dartrs
is a 🤗candle
backend inference library for Dart v2 models.
pip install -U dartrs
from dartrs.dartrs import DartTokenizer
from dartrs.utils import get_generation_config
from dartrs.v2 import (
compose_prompt,
MistralModel,
V2Model,
)
import time
import os
MODEL_NAME = "p1atdev/dart-v2-base"
model = MistralModel.from_pretrained(MODEL_NAME)
tokenizer = DartTokenizer.from_pretrained(MODEL_NAME)
config = get_generation_config(
prompt=compose_prompt(
copyright="vocaloid",
character="hatsune miku",
rating="general", # sfw, general, sensitive, nsfw, questionable, explicit
aspect_ratio="tall", # ultra_wide, wide, square, tall, ultra_tall
length="medium", # very_short, short, medium, long, very_long
prompt="1girl, cat ears",
do_completion=False
),
tokenizer=tokenizer,
)
start = time.time()
output = model.generate(config)
end = time.time()
print(output)
print(f"Time taken: {end - start:.2f}s")
# cowboy shot, detached sleeves, empty eyes, green eyes, green hair, green necktie, hair in own mouth, hair ornament, letterboxed, light frown, long hair, long sleeves, looking to the side, necktie, parted lips, shirt, sleeveless, sleeveless shirt, twintails, wing collar
# Time taken: 0.26s
Prompt Format
prompt = (
f"<|bos|>"
f"<copyright>{copyright_tags_here}</copyright>"
f"<character>{character_tags_here}</character>"
f"<|rating:general|><|aspect_ratio:tall|><|length:long|>"
f"<general>{general_tags_here}"
)
Rating tag:
<|rating:sfw|>
,<|rating:general|>
,<|rating:sensitive|>
,nsfw
,<|rating:questionable|>
,<|rating:explicit|>
sfw
: randomly generates tags ingeneral
orsensitive
rating categories.general
: generates tags ingeneral
rating category.sensitive
: generates tags insensitive
rating category.nsfw
: randomly generates tags inquestionable
orexplicit
rating categories.questionable
: generates tags inquestionable
rating category.explicit
: generates tags inexplicit
rating category.
Aspect ratio tag:
<|aspect_ratio:ultra_wide|>
,<|aspect_ratio:wide|>
,<|aspect_ratio:square|>
,<|aspect_ratio:tall|>
,<|aspect_ratio:ultra_tall|>
ultra_wide
: generates tags suits for extremely wide aspect ratio images. (~2:1)wide
: generates tags suits for wide aspect ratio images. (2:1~9:8)square
: generates tags suits for square aspect ratio images. (9:8~8:9)tall
: generates tags suits for tall aspect ratio images. (8:9~1:2)ultra_tall
: generates tags suits for extremely tall aspect ratio images. (1:2~)
Length tag:
<|length:very_short|>
,<|length:short|>
,<|length:medium|>
,<|length:long|>
,<|length:very_long|>
very_short
: totally generates ~10 number of tags.short
: totally generates ~20 number of tags.medium
: totally generates ~30 number of tags.long
: totally generates ~40 number of tags.very_long
: totally generates 40~ number of tags.
Model Details
Model Description
- Developed by: Plat
- Model type: Causal language model
- Language(s) (NLP): Danbooru tags
- License: Apache-2.0
- Demo: Available on 🤗 Space
Training Details
Training Data
This model was trained with:
- isek-ai/danbooru-tags-2024 with revision
202403-at20240423
: 7M size of danbooru tags dataset since 2005 to 2024/03/31.
Training Procedure
TODO
Preprocessing [optional]
[More Information Needed]
Training Hyperparameters
The following hyperparameters were used during training:
- learning_rate: 0.0005
- train_batch_size: 1024
- eval_batch_size: 256
- seed: 42
- gradient_accumulation_steps: 2
- total_train_batch_size: 2048
- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
- lr_scheduler_type: cosine
- lr_scheduler_warmup_steps: 1000
- num_epochs: 5
Evaluation
Evaluation has not been done yet and it needs to evaluate.
Model Architecture and Objective
The architecture of this model is Mistral. See details in config.json.
Compute Infrastructure
Private server.
Hardware
8x RTX A6000
Software
- Dataset processing: 🤗 Datasets
- Training: 🤗 Transformers
- SFT: 🤗 TRL
- Inference library: 📦 dartrs
- Backend: 🤗 candle
Related Projects
- dart-v1: The first version of the Dart model.
- KBlueLeaf/DanTagGen: The Aspect Ratio tag was inspired by this project.
- furusu/danbooru-tag-similarity: The idea of clustering tags and its training method was inspired by this project.