LoRA로 StableDiffusion 그림체 학습시키기

이번 글에서는 LoRA를 이용하여 StableDiffusion을 FineTuning하는 법에 대해 포스팅하고자 한다!

LoRA

Low-Rank Adaptation of Large Language Models(LoRA)는 pretrained 모델의 weights를 고정하고, 그 위에 새로운 task를 위한 weights를 추가하는 방식으로 fine tuning을 진행한다.

이 방식은 Trainable한 parameter의 수를 크게 줄여, 기존의 fine tuning 방식보다 훨씬 빠르다.

또한, 기존의 fine tuning 방식은 모든 weights를 업데이트하는데 반해, LoRA는 일부 weights만 업데이트하기 때문에, 기존의 fine tuning 방식보다 더 적은 양의 데이터로 fine tuning을 할 수 있다.

LoRA

 

LoRA를 이용하면, 각 task 별로, pretrained 모델에 얹을 task specific한 파라미터 조금만 저장 및 로드하면 된다.

다른 말로, pretrained 모델은 공유하고, 작은 LoRA module을 task 별로 빌드할 수 있다.

이를 통해 용량과 스위칭에 드는 비용을 아낄 수 있다. 대부분의 파라미터에 대해서 gradient 계산이나 최적화가 필요없고,

inject한 적은 양의 low rank matrices만 최적화해주면된다.

Trainable parameters 확인용

 

환경 준비

하드웨어 성능이 안 따라주기 때문에 Colab Pro을 사용하였다🤣

데이터셋

나는 OilCanvas느낌의, 즉 유화 느낌의 그림체로 학습시키고 싶었기 때문에 유화 데이터셋을 수집해서 커스텀을 진행하였다. 데이터는 캐글에서 수집을 하였다.

허깅페이스에서 수집한 데이터에 대해서 BLIP-2모델을 활용하여 캡션을 생성하였고, 최종 데이터셋을 완성했다.

BLIP-2을 활용한 Image Caption생성은 https://huggingface.co/blog/blip-2 여기서 소개한 사용법을 그대로 사용했다.

 

우선, 유화느낌의 데이터셋만 필요하기 때문에 impressionism 장르만 추출하였다.

#사진 저장 (impressionism만)
path='/content/drive/MyDrive/data/image/test2'
from PIL import Image
for i in range(len(dataset['train'])):
  if ((dataset['train'][i]['style'][0])==12): #impressionism
    gc.collect()
    torch.cuda.empty_cache()
    img=dataset['train'][i]['image'][0].save(path+'/impressionism_{}.jpg'.format(i),'jpeg') #이미지 저장
    gc.collect()
    torch.cuda.empty_cache()

 

위 코드에서 path는 자기 Drive경로로 지정해주면 된다.

 

그 후, blip모델을 이용해서 caption을 생성해주면 되는데, 아래 코드를 통해 수행할 수 있다. 마찬가지로 '/content/drive/MyDrive/data/image/test2' 이 부분은 위에서 지정한 path로 설정해주면 된다.

import os
import pandas as pd

prompt = "a painting of"
caption=[]
img_name=[]

file_list=os.listdir('/content/drive/MyDrive/data/image/test2')

for i in range(len(file_list)):
  gc.collect()
  torch.cuda.empty_cache()

  #이미지 이름 저장
  img_name.append(file_list[i])

  #이미지 업로드
  img=Image.open('/content/drive/MyDrive/data/image/test2/'+file_list[i]).convert('RGB')
  img=img.resize((256,256))

  gc.collect()
  torch.cuda.empty_cache()

  inputs = processor(img, text=prompt, return_tensors="pt").to(device, torch.float16)
  generated_ids = model.generate(**inputs, max_new_tokens=600)
  generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
  caption.append(generated_text)

  gc.collect()
  torch.cuda.empty_cache()

 

위 코드를 돌리면 아래와 같이 각 이미지에 대해 caption이 생성되는 것을 확인할 수 있을 것이다.

허깅페이스에 올려두었으니 해당 과정(이미지 캡션 생성)은 생략하고 이미지만 저장해도 된다.

파인 튜닝

파인튜닝하는 코드는 https://github.com/huggingface/diffusers/tree/main/examples/text_to_image 여기서 소개한 사용법을 거의 그대로 사용했다. Text-to-Image Diffusion 파인 튜닝을 위한 LoRA 구현이다.

pretrained model은 runwayml/stable-diffusion-v1-5을 사용했다.

 

이때, 이미지 경로 에러가 발생할 수도 있는데, 아래처럼 train_text_to_img_lora.py파일을 수정해주면 된다.

def preprocess_train(examples):
	# # 수정한 코드
    images = [Image.open('/content/drive/MyDrive/data/image/test2/'+image).convert("RGB") for image in examples[image_column]]
    # # 원래 코드
    #images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["input_ids"] = tokenize_captions(examples)
    return examples

 

이미지 경로를 직접 지정을 해주어 에러를 해결할 수 있다. 이때, /content/drive/MyDrive/data/image/test2/ 이 부분은 각자 개인 컴퓨터에 맞는 드라이브 경로로 지정해줘야 한다. (해당 경로에 이미지 데이터가 저장이 된 상태여야 한다. 위에서 말한 path와 동일한 경로)

 

수정 한 후, 아래 코드를 통해 FineTuning하면 된다.

!CUDA_VISIBLE_DEVICES="0" accelerate launch train_text_to_image_lora.py \
  --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
  --dataset_name="Eunju2834/oil_impressionism_style" --caption_column="text" \
  --resolution=512 --random_flip \
  --train_batch_size=1 \
  --num_train_epochs=20 --checkpointing_steps=10000 \
  --learning_rate=1e-04 --lr_scheduler="cosine" --lr_warmup_steps=0 \
  --seed=2024 \
  --output_dir="/content/drive/MyDrive/sd_model_fintuning_w.LoRA/fintune model" \
  --validation_prompt="Oil Painting, Impressionism, oil painting with brushstrokes, bichu, Park stroll, joyful atmosphere, laughter-filled time, playful dogs, vibrant park scene, cheerful interactions, happy pet owners, heartwarming moments, vibrant community vibes "  --report_to="wandb"

위 코드를 실행하면 각 checkpoint별로 .safetensor 파일이 output_dir 경로로 떨어진다.

safetensor는 tensor를 저장하는 포맷이다. 즉, matrices, 즉, weights를 저장한다. 이걸 그대로 사용할 순 없고 pretrained model에 적용해서 써야한다. 그 방법은 아래와 같다.

결과 weights 사용하기

모델 weights는 https://huggingface.co/Eunju2834/LoRA_oilcanvas_style  허깅페이스에 올려두었다.

from diffusers import StableDiffusionPipeline
import torch
import os

model_path='Eunju2834/LoRA_oilcanvas_style' #hugging face

pipe = StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16,use_auth_token=True)
pipe.unet.load_attn_procs(model_path)
pipe.to("cuda")

#positive prompt
prompt = ''' (Oil Painting : 1.1), (Impressionism : 1.2) ,(oil painting with brush strokes: 1.2),
Park stroll, joyful atmosphere, laughter-filled time, playful dogs, vibrant park scene,
cheerful interactions, happy pet owners, heartwarming moments, vibrant community vibes'''

#negative prompt
neg_prompt='''FastNegativeV2,(bad-artist:1.0),
(worst quality, low quality:1.4), (bad_prompt_version2:0.8),
bad-hands-5,lowres, bad anatomy, bad hands, ((text)), (watermark),
error, missing fingers, extra digit, fewer digits, cropped,
worst quality, low quality, normal quality, ((username)), blurry,
 (extra limbs), bad-artist-anime, badhandv4, EasyNegative,
 ng_deepnegative_v1_75t, verybadimagenegative_v1.3, BadDream,
(three hands:1.1),(three legs:1.1),(more than two hands:1.4),
(more than two legs,:1.2),badhandv4,EasyNegative,ng_deepnegative_v1_75t,verybadimagenegative_v1.3,(worst quality, low quality:1.4),text,words,logo,watermark,
'''



image = pipe(prompt, negative_prompt=neg_prompt,num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("/content/drive/MyDrive/sd_model_fintuning_w.LoRA/image/oil_impressionism_park_stroll2.png")

 

prompt에는 이미지 생성 프롬프트를, neg_prompt에는 negative prompt를 넣어주어 프롬프트 엔지니어링을 하면 된다. 

 

위 프롬프트에 대한 결과물은 아래와 같다.

 

결과물을 보면 확실히 유화느낌이 강한 것을 확인할 수 있다.