Stable Diffusion 모델과 Fine-Tuning 방법 - LoRA

Stable Diffusion 모델과 Fine-Tuning 방법

Stable Diffusion (SD) 모델은 텍스트 프롬프트를 기반으로 고품질 이미지를 생성하는 딥러닝 기반 확산 모델입니다. 주로 Latent Diffusion이라는 기술을 바탕으로 하며, 텍스트에서 이미지로의 매핑을 학습한 후 Latent Space(잠재 공간)에서 이미지 생성 작업을 수행합니다. 이를 통해 높은 해상도와 디테일을 유지하면서도 효율적인 이미지 생성을 가능하게 합니다.

 

Stable Diffusion은 미술, 일러스트, 사진 생성과 같은 다양한 분야에서 사용되며, 이를 보다 맞춤화된 방식으로 활용하기 위해 Fine-Tuning (파인 튜닝)이 종종 필요합니다.

---

Stable Diffusion 파인 튜닝이란?

Fine-Tuning은 기존에 훈련된 Stable Diffusion 모델에 새로운 데이터나 원하는 특성을 추가로 학습시키는 과정입니다. 이를 통해 특정 스타일, 인물, 객체를 모델이 더 잘 이해하고 생성할 수 있도록 개선합니다.

파인 튜닝의 주요 방법들

1. DreamBooth
   • 특정 객체나 인물을 모델이 기억하고 생성할 수 있도록 훈련합니다.
   • 소량의 이미지만으로 특정 객체를 모델에 주입할 수 있어 편리합니다.
2. LoRA (Low-Rank Adaptation)
   • 모델의 일부 파라미터만 조정하여 훈련 속도를 높이고 메모리 사용량을 줄입니다.
   • 원본 모델의 성능을 유지하면서도 다양한 스타일을 적용할 수 있습니다.
3. Textual Inversion
   • 모델의 가중치를 직접 조정하지 않고 새로운 임베딩(embedding) 토큰을 학습합니다.
   • 모델이 새로운 개념이나 스타일을 이해할 수 있도록 텍스트 프롬프트를 확장합니다.

2024.12.08 - [AI] - Stable Diffusion 모델 파인 튜닝 가이드

Stable Diffusion 모델 파인 튜닝 가이드

 

Stable Diffusion 모델을 LoRA를 이용해 파인 튜닝하는 방법

최근 Stable Diffusion은 고품질 이미지를 생성할 수 있는 강력한 모델로 주목받고 있습니다. 그러나 기본 모델만으로는 특정 스타일이나 주제에 적합한 이미지를 생성하는 데 제한이 있을 수 있습니다. 이를 해결하기 위해 Low-Rank Adaptation (LoRA) 를 이용한 파인 튜닝 방법이 인기를 끌고 있습니다. 이번 글에서는 Stable Diffusion 모델을 LoRA로 튜닝하는 방법을 단계별로 소개합니다.

 

LoRA란 무엇인가?

LoRA는 파인 튜닝 과정에서 전체 모델 파라미터를 업데이트하지 않고, 저차원 행렬(rank decomposition) 을 추가 학습시키는 방식입니다.

• 장점:
  • 메모리 효율적: 기존 방식보다 VRAM 사용량이 적습니다.
  • 속도: 학습 속도가 빠릅니다.
  • 확장성: 하나의 베이스 모델로 다양한 LoRA 가중치를 적용하여 여러 스타일을 사용할 수 있습니다.

준비물

LoRA를 사용해 Stable Diffusion 모델을 파인 튜닝하려면 다음이 필요합니다:

• Stable Diffusion 기본 모델: 예를 들어, SD 1.5 또는 SDXL.
• 파인 튜닝용 데이터셋: 특정 스타일이나 주제를 반영하는 이미지와 캡션.
• 파인 튜닝 도구: diffusers 라이브러리와 LoRA 지원 툴킷.
• GPU: NVIDIA GPU(CUDA 지원) 권장.

 

환경 설정

Python 환경과 필요한 라이브러리를 설치합니다.

# Python 환경 설정
python -m venv lora-env
source lora-env/bin/activate

# 필요한 라이브러리 설치
pip install torch torchvision transformers diffusers accelerate
pip install datasets safetensors

 

데이터셋 준비

파인 튜닝을 위해 적절한 이미지와 텍스트 캡션이 포함된 데이터셋을 준비합니다.

• 이미지 포맷: JPG/PNG 등 일반적인 포맷.
• 텍스트 캡션: 각 이미지의 내용을 설명하는 간단한 문장.

이미지와 캡션은 다음과 같은 구조로 준비합니다:

/dataset/
    /images/
        img1.jpg
        img2.jpg
    captions.txt  # "img1.jpg: A cat sitting on a chair"

 

LoRA 파인 튜닝 코드

다음은 LoRA를 이용해 Stable Diffusion 모델을 파인 튜닝하는 예제 코드입니다.

from diffusers import StableDiffusionPipeline
from diffusers.models.lora import LoRA
import torch

# 1. Stable Diffusion 모델 불러오기
model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

# 2. LoRA 구성
lora = LoRA(model=pipe.unet, rank=4)  # Rank는 조정 가능 (작을수록 메모리 효율적)

# 3. 데이터 로딩
from datasets import load_dataset

dataset = load_dataset("path/to/your/dataset", split="train")

# 4. 파인 튜닝 설정
from torch.optim import AdamW

optimizer = AdamW(lora.parameters(), lr=5e-5)

# 5. 학습 루프
epochs = 3
for epoch in range(epochs):
    for batch in dataset:
        images, captions = batch["images"], batch["captions"]
        latents = pipe.vae.encode(images).latent_dist.sample()  # 이미지 인코딩
        noise = torch.randn_like(latents).to("cuda")
        
        # 모델 예측 및 손실 계산
        predicted_noise = pipe.unet(latents, noise, captions).sample
        loss = ((predicted_noise - noise) ** 2).mean()

        # 역전파 및 업데이트
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    print(f"Epoch {epoch + 1}/{epochs}, Loss: {loss.item()}")

# 6. LoRA 저장
lora.save_pretrained("path/to/save/lora")

 

LoRA 적용

학습이 완료된 LoRA를 Stable Diffusion 모델에 적용하는 방법은 다음과 같습니다.

from diffusers.models.lora import LoRA

# LoRA 가중치 불러오기
lora = LoRA.from_pretrained("path/to/save/lora")

# LoRA를 모델에 추가
pipe.unet = lora.apply(pipe.unet)

# 테스트
prompt = "A fantasy landscape with dragons"
image = pipe(prompt).images[0]
image.show()

 

결과 및 활용

LoRA를 사용하면 Stable Diffusion 기본 모델을 수정하지 않고도 원하는 스타일이나 주제를 가진 이미지를 생성할 수 있습니다. 예를 들어:

• 특정 예술 스타일
• 특정 인물이나 캐릭터의 일러스트
• 기업 또는 프로젝트의 브랜드 이미지

 

추가 팁

• Rank 값 조정: Rank 값을 늘리면 더 복잡한 스타일을 학습할 수 있지만, 메모리 사용량이 증가합니다.
• 데이터 증강: 데이터셋 크기가 작다면 이미지 회전, 색상 조정 등 데이터 증강을 고려하세요.
• 커뮤니티 활용: Hugging Face에서 LoRA 가중치를 공유하고 커뮤니티의 작업을 참고하세요.

 

LoRA를 활용한 Stable Diffusion 파인 튜닝은 적은 자원으로도 효과적인 결과를 낼 수 있는 방법입니다. 여러분의 창작 프로젝트에 이 기술이 큰 도움이 되길 바랍니다! 🎨