1、使用diffusers-0.14.0, stabel-diffusion 模型 v-1.5版本
下载diffusers-0.14.0 并解压,新建文件test.py, 写入以下:
import torch
from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained("./stable-diffusion-v1-5")
pipe = pipe.to("cuda")
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
执行以上过程,自动下载新建stable-diffusion-v1-5文件夹并下载多个模型在次目录下,
并根据prompt 生成一张image。
2、自己数据集的制作。
使用X-decoder 项目的demo_captioning.py模型,https://github.com/microsoft/X-Decoder
每一张图片生成一个对应的txt文件,txt文件中保存了该图片的prompt. demo_captioning.py 需要小改一下保存prompt就行。
将所有图片和txt都存放到同一个文件夹下。
因为模型训练使用的是huggingFace 的datasets 格式,所有再生成一个csv文件,或者json文件都行。代码如下:
import csv
import glob
header = ['file_name','text']
img_list = glob.glob('./172_186/prompt_imgs/*.jpg')
with open('metadata.csv', 'w', encoding='utf-8') as file_obj:
# 1:创建writer对象
writer = csv.writer(file_obj)
# 2:写表头
writer.writerow(header)
# 3:遍历列表,将每一行的数据写入csv
for p in img_list:
txt = p.replace('jpg','txt')
content = [p.split('/')[-1],txt]
writer.writerow(content)
测试一下生成csv文件格式对不对:
data_files = {}
data_files["train"] = os.path.join('./pokman/prompts', "**")
dataset = load_dataset(
"imagefolder",
data_files=data_files)
print(dataset['train'])
column_names = dataset["train"].column_names
print(column_names)
输出结果:
Dataset imagefolder downloaded and prepared to /home/lyn/.cache/huggingface/datasets/imagefolder/default-a6d71509cd596a17/0.0.0/37fbb85cc714a338bea574ac6c7d0b5be5aff46c1862c1989b20e0771199e93f. Subsequent calls will reuse this data.
Dataset({
features: ['image', 'text'],
num_rows: 2151
})
['image', 'text']
最后把这个csv文件一起放入Image和txt的文件夹中即可。
3、模型训练。
脚本diffusers-0.14.0/examples/text_to_image/train_text_to_image_lora.py
命令:
先设置以下环境变量,执行
export MODEL_NAME="/diffusers-0.14.0/stable-diffusion-v1-5" 模型所在目录
export OUTPUT_DIR="/diffusers-0.14.0/logs" 模型文件保存目录
export IMAGE_FOLDER='/172_186/prompt_imgs' 训练数据所在目录
然后终端执行:
accelerate launch --mixed_precision="fp16" examples/text_to_image/train_text_to_image_lora.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataloader_num_workers=8 \
--resolution=512 --center_crop --random_flip \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--max_train_steps=15000 \
--learning_rate=1e-04 \
--max_grad_norm=1 \
--lr_scheduler="cosine" --lr_warmup_steps=0 \
--output_dir=${OUTPUT_DIR} \
--report_to=wandb \
--checkpointing_steps=500 \
--seed=1337 \
--train_data_dir=${IMAGE_FOLDER} \