LoRA-水墨风图片生成

★★★ 本文源自AlStudio社区精品项目，【点击此处】查看更多精品内容 >>>

1、安装依赖

环境要求是用V100或A100，并保证显存有32GB，不然训练会爆显存

注意！！！安装完依赖，记得点一下上方得“重启内核”按钮，否则后面还是会找不到 ppdiffusers 模块

!pip install paddlenlp --user
!pip install ppdiffusers==0.11.1 --user

2、解压水墨风数据集

import os
if not os.path.exists('img'):
    os.mkdir('img')
if not os.path.exists('img/painting'):
    os.system('unzip data/data180682/painting.zip -d img/ink_cn')
# if not os.path.exists('img/2d_head'):
#     os.system('unzip data/data183241/extra_data.zip -d img/2d_head')
# if not os.path.exists('img/2d_person'):
#     os.system('unzip data/data159362/gallary.zip -d img/2d_person')

3、加载初始模型 （初次5分钟）

import os
# 如果之前下载过模型，就加载一下，防止重新下载
if not os.path.exists('.paddlenlp'):
    os.mkdir('.paddlenlp')
    !cp -r work/* /home/aistudio/.paddlenlp/

from ppdiffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import paddle
import PIL
from random import random

# 初次下载模型，大约5分钟
pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

# 备份模型文件，防止下次打开重新下载，大约1分钟
# !cp -r /home/aistudio/.paddlenlp/* /home/aistudio/work/

# # 原始模型保存为LoRA模型
# from ppdiffusers.models.cross_attention import LoRACrossAttnProcessor
# from ppdiffusers.modeling_utils import freeze_params, unwrap_model
# lora_attn_procs = {}
# for name in unet.attn_processors.keys():
#     cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
#     if name.startswith("mid_block"):
#         hidden_size = unet.config.block_out_channels[-1]
#     elif name.startswith("up_blocks"):
#         block_id = int(name[len("up_blocks.")])
#         hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
#     elif name.startswith("down_blocks"):
#         block_id = int(name[len("down_blocks.")])
#         hidden_size = unet.config.block_out_channels[block_id]

#     lora_attn_procs[name] = LoRACrossAttnProcessor(
#         hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=4
#     )
# # freeze_params(unet.parameters())
# pipe.unet.set_attn_processor(lora_attn_procs)
# unwrap_model(pipe.unet).save_attn_procs("src")

4、查看初始效果

初始模型生成的会偏照片一点

import os
def go(prompt, seed=42, steps=10):
    paddle.seed(seed)
    image = pipe(prompt, num_inference_steps=steps).images[0]
    name = prompt[:20]
    if not os.path.exists('result'):
        os.mkdir('result')
    print('prompt:', prompt)
    print('seed:', seed, 'steps:', steps)
    image.save(f"result/se{seed}_st{steps}_{name}_demo2.png")
    image.show()

go('a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.', seed=7434, steps=15)

  0%|          | 0/15 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 7434 steps: 15

go('mountain', seed=6, steps=20)

  0%|          | 0/20 [00:00<?, ?it/s]


prompt: mountain
seed: 6 steps: 20

go('fuchun river', seed=8, steps=30)

  0%|          | 0/30 [00:00<?, ?it/s]


prompt: fuchun river
seed: 8 steps: 30

5、训练模型 （10分钟）

这一步可以改数据集，把 --instance_data_dir 这个参数对应的值改为包含待训练图片的文件夹目录即可

import os
if not os.path.exists('out'):
    print('开始训练LoRA模型，大约10分钟')
    instance_data_dir = '/home/aistudio/img/ink_cn/painting/trainB/'
    # instance_data_dir = '/home/aistudio/img/2d_head/extra_data/images/'
    # instance_data_dir = '/home/aistudio/img/2d_person/gallary/'
    num_train_epochs = 1
    os.system(f'python /home/aistudio/train_dreambooth_lora.py --num_train_epochs="{num_train_epochs}" --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5"    --instance_data_dir="{instance_data_dir}"   --output_dir="/home/aistudio/out/"   --instance_prompt="a photo of sks dog"   --resolution=512   --train_batch_size=1   --gradient_accumulation_steps=1   --checkpointing_steps=100   --learning_rate=1e-4   --report_to="visualdl"   --lr_scheduler="constant"   --lr_warmup_steps=0   --max_train_steps=500   --validation_prompt="A photo of sks dog in a bucket"   --validation_epochs=50   --lora_rank=4   --seed=0')
else:
    print('已训练过模型')

已训练过模型

加载一下训练好的lora

# 加载训练好的lora
pipe.unet.load_attn_procs("out", from_hf_hub=True)

6、生成水墨风图片~

steps为100时：V100 20秒1张，A100 10秒1张，结果会保存在 result 文件夹下

参数解释：

prompt: 提示词，代表你想要AI生成的东西

seed: 种子，固定的 prompt+seed 可以生成固定的图片

steps: 迭代次数，次数越多耗时越长，效果越好，一般建议在10~50之间

先跑几张步数少的，看到风格还喜欢的，就增加到步数，输入对应 种子+步数 来精修

prompt = 'a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.'
for i in range(10):
    seed = int(random()*10000)
    go(prompt, seed)

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 1104 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 1322 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 7157 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 994 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 460 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 669 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 689 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 5119 steps: 10


  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 8500 steps: 10

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 3466 steps: 10

7、最终效果

个人试了下，小狗、山水画效果还是不错的，欢迎fork尝试画其他水墨画😜

（画人有点恐怖，谨慎尝试；如果prompt学得好，还能画出来全黑的图😂）

go(prompt, seed=7434, steps=15)

  0%|          | 0/15 [00:00<?, ?it/s]


prompt: a portrait of shiba inu with a green cap growing on its head. intricate. lifelike. soft light.
seed: 7434 steps: 15

go('mountain', seed=6, steps=20)

  0%|          | 0/20 [00:00<?, ?it/s]


prompt: mountain
seed: 6 steps: 20

go('fuchun river', seed=8, steps=10)

  0%|          | 0/10 [00:00<?, ?it/s]


prompt: fuchun river
seed: 8 steps: 10

hun river
seed: 8 steps: 10

此文章为搬运
原项目链接