Dart V2 Moe Sft

🚀 Dart (Danbooru Tags Transformer) v2

Dart (Danbooru Tags Transformer) v2 是一个经过微调的模型，专门用于生成 Danbooru 标签。它能根据输入的相关信息，准确且高效地生成对应的标签，为相关应用场景提供了有力支持。

Demo: 🤗 Space with ZERO

✨ 主要特性

模型变体

📦 安装指南

使用 📦dartrs 库

pip install -U dartrs

💻 使用示例

基础用法

使用 🤗Transformers

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"<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

高级用法

使用 📦dartrs 库

⚠️ 重要提示

这个库目前还处于实验阶段，未来可能会有重大变更。

📦dartrs 是一个基于 🤗candle 后端的 Dart v2 模型推理库。

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 = (
    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:sfw|>、<|rating:general|>、<|rating:sensitive|>、nsfw、<|rating:questionable|>、<|rating:explicit|>
  • sfw：随机生成 general 或 sensitive 评级类别的标签。
  • general：生成 general 评级类别的标签。
  • sensitive：生成 sensitive 评级类别的标签。
  • nsfw：随机生成 questionable 或 explicit 评级类别的标签。
  • questionable：生成 questionable 评级类别的标签。
  • explicit：生成 explicit 评级类别的标签。
• 宽高比标签：<|aspect_ratio:ultra_wide|>、<|aspect_ratio:wide|>、<|aspect_ratio:square|>、<|aspect_ratio:tall|>、<|aspect_ratio:ultra_tall|>
  • ultra_wide：生成适合极宽宽高比图像（约 2:1）的标签。
  • wide：生成适合宽宽高比图像（2:1 - 9:8）的标签。
  • square：生成适合正方形宽高比图像（9:8 - 8:9）的标签。
  • tall：生成适合高宽高比图像（8:9 - 1:2）的标签。
  • ultra_tall：生成适合极高宽高比图像（1:2 及以上）的标签。
• 长度标签：<|length:very_short|>、<|length:short|>、<|length:medium|>、<|length:long|>、<|length:very_long|>
  • very_short：总共生成约 10 个标签。
  • short：总共生成约 20 个标签。
  • medium：总共生成约 30 个标签。
  • long：总共生成约 40 个标签。
  • very_long：总共生成 40 个以上的标签。
• 身份标签：<|identity:none|>、<|identity:lax|>、<|identity:strict|>
  • 此标签指定了在生成标签时，对所提供标签中角色或主体身份的保留严格程度。
  • none：当指定的通用标签非常少的时候推荐使用。它会非常有创造性地生成标签，但有时会忽略通用标签的条件。
  • lax：如果您希望保留通用标签中角色或主体的身份，推荐使用此标签。它会尽量不生成与输入通用标签冲突的标签。
  • strict：如果您强烈希望保留通用标签中角色或主体的身份，推荐使用此标签。它会比 lax 更严格地避免生成与输入通用标签冲突的标签，但创造性较差。如果您不喜欢 strict 的结果，请尝试 lax 或 none。

模型详情

模型描述

训练详情

训练数据

该模型使用以下数据进行训练：

• isek-ai/danbooru-tags-2024，版本为 202403-at20240423：这是一个从 2005 年到 2024 年 3 月 31 日的 Danbooru 标签数据集，大小为 700 万。

训练过程

TODO

预处理 [可选]

[需要更多信息]

训练超参数

训练过程中使用了以下超参数：

• 学习率：0.00025
• 训练批次大小：1024
• 评估批次大小：256
• 随机种子：42
• 梯度累积步数：2
• 总训练批次大小：2048
• 优化器：Adam，β1 = 0.9，β2 = 0.999，ε = 1e-08
• 学习率调度器类型：余弦
• 学习率调度器预热步数：1000
• 训练轮数：4

评估

目前尚未进行评估，需要进行评估。

模型架构和目标

该模型的架构是 Mixtral。详细信息请参考 config.json。

计算基础设施

硬件

大学实验室的服务器，配备 8 块 RTX A6000 显卡。

软件

• 数据集处理：🤗 Datasets
• 训练：🤗 Transformers
• 监督微调：🤗 TRL
• 推理库：📦 dartrs
  • 后端：🤗 candle

相关项目

• dart-v1：Dart 模型的第一个版本。
• KBlueLeaf/DanTagGen：宽高比标签的灵感来源于此项目。
• furusu/danbooru-tag-similarity：标签聚类的想法及其训练方法的灵感来源于此项目。

📄 许可证

本模型使用 Apache-2.0 许可证。