为 Stable Diffusion 微调创建 Diffusers 兼容的数据集

为了吸引年轻人，我将用你们年轻人的行话说话。😎

哟，各位小伙伴们！🎉 今天，我们将来聊聊如何创建一个 diffusers 数据集（我知道，听起来有点无聊，但请坚持住！）。🤘

第一步：加载你的炫酷 Img2Img 流水线并将其发送到 GPU（不管那意味着什么）😎😎😎

from transformers import pipeline

pipe = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large", device="cuda") # I'm using blip, as that's what most SD models expect

第二步：休息一下（你#值得拥有！🏖️😎）

放松一下，我们再集合。去拿点零食什么的。🍕🌟

第三步：很高兴你回来，现在让我们设置你的超酷图片文件夹 📁✨

import os

folder = "your_image_folder" # file structure is your_image_folder/train/*.png
image_folder = os.path.join(folder, "train")

第四步：太累了，复制粘贴了这么多字符，还没署名，再休息一下吧🥵🥵🥵🥵🔥

让你的手指休息一下，因为我们很快就要继续了。🤯💥

第……我忘了……第 4 步？我想？像 #老板 一样加载那个文件夹 🚀🕶️

from datasets import load_dataset

dataset = load_dataset('imagefolder', data_dir=image_folder, split='train')

第 4 步（没错，听起来对）。有型地制作你的格式化函数 🧙‍♂️🎨

def generate_captions(examples):
    result = pipe(examples['image'])
    captions = [item[0]['generated_text'] for item in result]
    return {"text": captions}

第 4 步。映射数据集（又数据？真无聊）

dataset = dataset.map(generate_captions, batched=True, batch_size=8) # default batch size is 1000

如果你的电脑在这之后着火了 🔥 或爆炸了 💥，别怪我 😎💪。

第 4 步。把训练集找回来 😎😎🔪

from datasets import DatasetDict

dataset_dict = DatasetDict({
    'train': dataset
})

dataset_dict # test to see if the new columns are there in train, if you were paying attention they should be

第 4 步。推送到 Hub（现在是个酷炫的俱乐部了吗？🤷‍♂️）💥🌈

dataset_dict.push_to_hub("your_output_repo")

第 4 步。🎉 你做到了！🎉 (终于)

现在去看 Vine 吧，或者别的什么。等等——不对……你们年轻人用的是那个新的短视频平台……Tik Tak？