dart-v2-base / README.md
p1atdev's picture
Update README.md
588265d verified
metadata
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 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.

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:

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

Related Projects