import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "p1atdev/dart-v2-sft"

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, cat ears<|identity:none|><|input_end|>"
)
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() != ""]))
# vocaloid, hatsune miku, 1girl, cat ears, closed mouth, detached sleeves, dress, expressionless, from behind, full body, green theme, hair ornament, hair ribbon, headphones, high heels, holding, holding microphone, long hair, microphone, monochrome, necktie, ribbon, short dress, shoulder tattoo, simple background, sleeveless, sleeveless dress, spot color, standing, tattoo, thighhighs, twintails, very long hair, white background

Using 📦`dartrs` library

[!WARNING] 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-sft"

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
        identity="none", # none, lax, strict
        prompt="1girl, cat ears",
    ),
    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}<|identity:none|><|input_end|>"
)

Rating tag: <|rating:sfw|>, <|rating:general|>, <|rating:sensitive|>, nsfw, <|rating:questionable|>, <|rating:explicit|>
- sfw: randomly generates tags in general or sensitive rating categories.
- general: generates tags in general rating category.
- sensitive: generates tags in sensitive rating category.
- nsfw: randomly generates tags in questionable or explicit rating categories.
- questionable: generates tags in questionable rating category.
- explicit: generates tags in explicit 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.
Identity tag: <|identity:none|>, <|identity:lax|>, <|identity:strict|>
- This tag specifies how strictly to preserve identity of character or subject in provided tags.
- none: recommended if the specified general tags are very few. It generates tags very creatively, but sometimes ignores the condition of the general tags.
- lax: recommended if you want to keep the identity of charaacters or subjects in the general tags. This tag tries not to generate tags which conflict with the input general tags.
- strict: recommended if you strongly want to keep the identity of charaacters or subjects in the general tags. This tag tries not to generate tags which conflict with the input general tags more strictly than lax. But this is less creative, so if you don't like the result with strict, please try lax or none.

Model Details

Model Description

Developed by: Plat
Model type: Causal language model
Language(s) (NLP): Danbooru tags
License: Apache-2.0
Finetuned from model: dart-v2-moe-base
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.00025
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: 4

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

Server in a university laboratory

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.

p1atdev/dart-v2-sft

Dart (Danbooru Tags Transformer) v2

Model variants

Usage

Using 🤗Transformers

Using 📦dartrs library