🔥

Dart V2を使って全自動でイラストデータセットを作成する

2024/05/12に公開

昨日、Dartという良い感じにプロンプトを生成してくれるLLMのV2が発表されました。

今回はこれを使ってControlNet等の学習に必要なイラストを大量に自動生成する方法を紹介します。

Dartって何？

Danbooruタグをいい感じに生成・補完してくれるLLM（大規模言語モデル）です。

詳しくは作者様が解説してくださっているこちらの記事をご参照ください。

データセット自動生成

本題です。
今回はそんなに難しいことはせず、Dartでプロンプトを自動生成し、そのプロンプトをDiffusersに食わせて画像を生成します。

Dartの呼び出し

まずはプロンプト作成部分です。

def get_prompt(model):
    
    prompt = (
        f"<|bos|>"
        f"<copyright></copyright>"
        f"<character></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,
        )
    return ", ".join([tag for tag in tokenizer.batch_decode(outputs[0], skip_special_tokens=True) if tag.strip() != ""])

MODEL_NAME = "p1atdev/dart-v2-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16)

Dartを用いてプロンプトを生成する部分については関数化しておくと、あとで大量に画像を生成する際に呼び出しやすくなります。
この際、モデルをロードする部分まで関数に組み込むと無駄に処理時間が長くなってしまうので、外に出すように気を付けましょう。

特に注目すべきは以下の部分となります

prompt = (
    f"<|bos|>"
    f"<copyright></copyright>"
    f"<character></character>"
    f"<|rating:general|><|aspect_ratio:tall|><|length:long|>"
    f"<general>1girl"
)

<rating> ブロック: レーティングカテゴリ（nsfwなど）を表すブロック
<copyright> ブロック: 版権タグが入るブロック
<character> ブロック: キャラクター名を入れるブロック
<general> ブロック: 上記に該当しないプロンプトをいれるブロック

今回は、色々な用途につかえるデータセットを作りたかったため、<copyright>タグと <character>タグは空にしてプロンプトを生成します。
もし設定する場合は、<character>キャラクター名</character>のような形で指定してください。
また、キャラクター名などを入れたくない場合はブロックを消すのではなく、あくまでブロックに挟まれている文字列がない状態（<character></character>）にすることが大切です。

作者様によると、ブロックに何も挟まれていない状態は、キャラクター名が何も設定されていないことを明示的に表す効果もあるようなので、ブロック自体は残しておきましょう。

（使用しないブロックを消してもプロンプト生成は可能ですが、例えばキャラクターブロックを消してしまった場合は、生成されるプロンプトにランダムなキャラクター名が含まれるようになってしまいました
）

Generalは、作りたいデータの性質(すべてのデータで共通して持っていてもらいたい要素)に合わせて指定すると良いかと思います。
今回は人物ありきのイラストデータセットを作りたかったので1girlだけ指定しています。

画像生成部分

画像生成部分はいつも通りDiffusersのt2i pipelineを使います


def make_image(pipe, model):
    prompt = get_prompt(model)
    negative_prompt = "EasyNegative, (lip, lipstick, worst quality, low quality:1.4), (rib, monochrome:1.3), long torso, (bad_prompt_version2:0.8), (negative_hand-neg:0.8),"
    image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]
    return image, prompt


pipe = AutoPipelineForText2Image.from_pretrained(
    "cagliostrolab/animagine-xl-3.0", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")

pipe.load_textual_inversion(pretrained_model_name_or_path="./embbed/EasyNegative.safetensors", token="EasyNegative", local_files_only=True)
pipe.load_textual_inversion(pretrained_model_name_or_path="./embbed/bad_prompt_version2.pt", token="bad_prompt_version2", local_files_only=True)
pipe.load_textual_inversion(pretrained_model_name_or_path="./embbed/negative_hand-neg.pt", token="negative_hand-neg", local_files_only=True)

こちらも、画像生成部分だけ関数化し、モデルロード部分は関数の外に出しておきましょう。

また、DartではPositive Promptしか生成しないのでNegative Promptは自分で設定する必要があります。
DiffusersでもEmbeddingsが使えるので、以下の記事を参考にEmbeddingsも利用しています。

コード全量

最後に、Dartで生成したプロンプトをt2i pipelineに渡し、生成したい画像枚数分だけループさせます。
完成版がこちら

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from diffusers import AutoPipelineForText2Image

def get_prompt(model):
    
    prompt = (
        f"<|bos|>"
        f"<copyright></copyright>"
        f"<character></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,
        )
    return ", ".join([tag for tag in tokenizer.batch_decode(outputs[0], skip_special_tokens=True) if tag.strip() != ""])

def make_image(pipe, model):
    prompt = get_prompt(model)
    negative_prompt = "EasyNegative, (lip, lipstick, worst quality, low quality:1.4), (rib, monochrome:1.3), long torso, (bad_prompt_version2:0.8), (negative_hand-neg:0.8),"
    image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]
    return image, prompt

MODEL_NAME = "p1atdev/dart-v2-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16)

pipe = AutoPipelineForText2Image.from_pretrained(
    "cagliostrolab/animagine-xl-3.0", torch_dtype=torch.float16, use_safetensors=True
).to("cuda")

pipe.load_textual_inversion(pretrained_model_name_or_path="./embbed/EasyNegative.safetensors", token="EasyNegative", local_files_only=True)
pipe.load_textual_inversion(pretrained_model_name_or_path="./embbed/bad_prompt_version2.pt", token="bad_prompt_version2", local_files_only=True)
pipe.load_textual_inversion(pretrained_model_name_or_path="./embbed/negative_hand-neg.pt", token="negative_hand-neg", local_files_only=True)

for idx in range(10000):
    print(f"process: {idx} Start")
    image, prompt = make_image(pipe, model)
    image.save(f"./data/image/image_{idx}.png")
    with open(f"./data/caption/caption_{idx}.txt", mode='w') as f:
        f.write(prompt)

生成結果

コードを動かすとこのような感じで生成されます。

生成されたプロンプト

1girl, :t, bandaid, bandaid on arm, bandaid on face, bandaid on leg,
bare shoulders, black socks, blue shorts, blush, closed mouth, denim,
denim shorts, grey hair, hair between eyes, hair ornament, heart,
heart hair ornament, legs, long hair, micro shorts, off shoulder,
off-shoulder dress, one side up, pout, purple eyes, shirt, shirt tug,
shorts, simple background, socks, solo, standing, tachi-e, very long hair,
white background, white footwear, white shirt

生成されたイラスト