Stable diffusion LoRA 训练过程

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} \