--- library_name: transformers license: apache-2.0 datasets: - isek-ai/danbooru-tags-2024 base_model: p1atdev/dart-v2-moe-base tags: - trl - sft - optimum - danbooru inference: false --- # Dart (Danbooru Tags Transformer) v2 This model is a fine-tuned Dart (Danbooru Tags Transformer) model that generates danbooru tags. Demo: [🤗 Space with ZERO](https://huggingface.co/spaces/p1atdev/danbooru-tags-transformer-v2) ## Model variants |Name|Architecture|Param size|Type| |-|-|-|-| |[v2-moe-sft](https://huggingface.co/p1atdev/dart-v2-moe-sft)|Mixtral|166m|SFT| |[v2-moe-base](https://huggingface.co/p1atdev/dart-v2-moe-base)|Mixtral|166m|Pretrain| |[v2-sft](https://huggingface.co/p1atdev/dart-v2-sft)|Mistral|114m|SFT| |[v2-base](https://huggingface.co/p1atdev/dart-v2-base)|Mistral|114m|Pretrain| |[v2-vectors](https://huggingface.co/p1atdev/dart-v2-vectors)|Embedding|-|Tag Embedding| ## Usage ### Using 🤗Transformers ```py import torch from transformers import AutoTokenizer, AutoModelForCausalLM MODEL_NAME = "p1atdev/dart-v2-moe-sft" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16) prompt = ( f"<|bos|>" f"vocaloid" f"hatsune miku" f"<|rating:general|><|aspect_ratio:tall|><|length:long|>" f"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`](https://github.com/p1atdev/dartrs) is a [🤗`candle`](https://github.com/huggingface/candle) backend inference library for Dart v2 models. ```py pip install -U dartrs ``` ```py from dartrs.dartrs import DartTokenizer from dartrs.utils import get_generation_config from dartrs.v2 import ( compose_prompt, MixtralModel, V2Model, ) import time import os MODEL_NAME = "p1atdev/dart-v2-moe-sft" model = MixtralModel.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 ```py prompt = ( f"<|bos|>" f"{copyright_tags_here}" f"{character_tags_here}" f"<|rating:general|><|aspect_ratio:tall|><|length:long|>" f"{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](https://huggingface.co/p1atdev/dart-v2-moe-base) - **Demo:** Available on [🤗 Space](https://huggingface.co/spaces/p1atdev/danbooru-tags-transformer-v2) ## Training Details ### Training Data This model was trained with: - [isek-ai/danbooru-tags-2024](https://huggingface.co/datasets/isek-ai/danbooru-tags-2024/tree/202403-at20240423) 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 [Mixtral](https://huggingface.co/docs/transformers/model_doc/mixtral). See details in [config.json](./config.json). ### Compute Infrastructure Server in a university laboratory #### Hardware 8x RTX A6000 #### Software - Dataset processing: [🤗 Datasets](https://github.com/huggingface/datasets) - Training: [🤗 Transformers](https://github.com/huggingface/transformers) - SFT: [🤗 TRL](https://github.com/huggingface/trl) - Inference library: [📦 dartrs](https://github.com/p1atdev/dartrs) - Backend: [🤗 candle](https://github.com/huggingface/candle) ## Related Projects - [dart-v1](https://huggingface.co/p1atdev/dart-v1): The first version of the Dart model. - [KBlueLeaf/DanTagGen](https://huggingface.co/collections/KBlueLeaf/dantaggen-65f82fa9335881a67573556b): The Aspect Ratio tag was inspired by this project. - [furusu/danbooru-tag-similarity](https://huggingface.co/spaces/furusu/danbooru-tag-similarity): The idea of clustering tags and its training method was inspired by this project.