Diffusion 모델 파인튜닝: LoRA와 DreamBooth로 커스텀 이미지 생성

AI 기술

Diffusion 모델LoRADreamBooth파인튜닝이미지 생성

이 글은 누구를 위한 것인가

  • 특정 브랜드 스타일로 이미지를 생성하고 싶은 팀
  • DreamBooth로 특정 인물·제품의 이미지를 생성하려는 개발자
  • LoRA와 전체 파인튜닝의 차이를 이해하고 싶은 ML 엔지니어

들어가며

"우리 브랜드 마스코트가 다양한 상황에서 나오는 이미지가 필요하다." 매번 디자이너에게 의뢰하는 대신 LoRA 파인튜닝으로 특정 스타일을 학습시키면 수백 장의 이미지를 자동 생성할 수 있다.

이 글은 bluefoxdev.kr의 Diffusion 모델 활용 을 참고하여 작성했습니다.


1. 파인튜닝 방식 비교

[파인튜닝 방법 비교]

전체 파인튜닝 (Full Fine-tuning):
  - 모든 파라미터 업데이트
  - VRAM: 24GB+ 필요
  - 시간: 수 시간
  - 저장 용량: 수 GB
  - 용도: 완전히 다른 스타일 도메인

LoRA (Low-Rank Adaptation):
  - 저랭크 행렬만 학습 (1-4% 파라미터)
  - VRAM: 8-12GB
  - 시간: 30분-2시간
  - 저장 용량: 10-100MB
  - 용도: 스타일, 특정 캐릭터

DreamBooth:
  - 특정 개체를 고유 토큰으로 학습
  - 소수 이미지(3-30장)로 가능
  - "sks person" 같은 트리거 단어 사용
  - 용도: 특정 인물, 제품, 펫 개인화

Textual Inversion:
  - 새로운 임베딩 토큰만 학습
  - 가장 가볍고 빠름
  - 품질은 상대적으로 낮음

[학습 데이터 준비]
  이미지 수: LoRA 50-200장, DreamBooth 5-30장
  해상도: 512x512 또는 768x768
  다양성: 다양한 각도, 배경, 조명
  캡셔닝: BLIP/WD14로 자동 캡션 생성

2. LoRA 파인튜닝 구현

from diffusers import (
    StableDiffusionPipeline,
    DDPMScheduler,
    UNet2DConditionModel,
)
from peft import LoraConfig, get_peft_model
import torch
from pathlib import Path
from PIL import Image
import json

def prepare_training_data(image_dir: str, output_dir: str):
    """학습 데이터 준비: 이미지 리사이징 + 캡셔닝"""
    from transformers import BlipProcessor, BlipForConditionalGeneration
    
    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
    caption_model = BlipForConditionalGeneration.from_pretrained(
        "Salesforce/blip-image-captioning-base"
    )
    
    image_paths = list(Path(image_dir).glob("*.{jpg,png,jpeg}"))
    metadata = []
    
    for img_path in image_paths:
        img = Image.open(img_path).convert("RGB")
        img = img.resize((512, 512))
        
        # 자동 캡션 생성
        inputs = processor(img, return_tensors="pt")
        output = caption_model.generate(**inputs, max_new_tokens=50)
        caption = processor.decode(output[0], skip_special_tokens=True)
        
        # 트리거 단어 추가
        full_caption = f"sks style, {caption}"
        
        save_path = Path(output_dir) / img_path.name
        img.save(save_path)
        
        metadata.append({
            "file_name": img_path.name,
            "text": full_caption,
        })
    
    with open(Path(output_dir) / "metadata.jsonl", "w") as f:
        for item in metadata:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

def setup_lora_model(
    base_model: str = "runwayml/stable-diffusion-v1-5",
    rank: int = 16,
    alpha: int = 32,
) -> tuple:
    """LoRA 설정 및 모델 준비"""
    
    unet = UNet2DConditionModel.from_pretrained(base_model, subfolder="unet")
    
    lora_config = LoraConfig(
        r=rank,
        lora_alpha=alpha,
        init_lora_weights="gaussian",
        target_modules=["to_k", "to_q", "to_v", "to_out.0"],
    )
    
    unet = get_peft_model(unet, lora_config)
    unet.print_trainable_parameters()
    
    return unet

def generate_with_lora(
    prompt: str,
    lora_path: str,
    base_model: str = "runwayml/stable-diffusion-v1-5",
    num_images: int = 4,
) -> list[Image.Image]:
    """LoRA 가중치로 이미지 생성"""
    
    pipe = StableDiffusionPipeline.from_pretrained(
        base_model,
        torch_dtype=torch.float16,
    ).to("cuda")
    
    # LoRA 가중치 로드
    pipe.load_lora_weights(lora_path)
    
    images = pipe(
        prompt=prompt,
        num_images_per_prompt=num_images,
        num_inference_steps=30,
        guidance_scale=7.5,
    ).images
    
    return images

def merge_lora_to_base(lora_path: str, output_path: str, alpha: float = 1.0):
    """LoRA를 베이스 모델에 병합 (배포용)"""
    from diffusers import StableDiffusionPipeline
    
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float32,
    )
    
    pipe.load_lora_weights(lora_path)
    pipe.fuse_lora(lora_scale=alpha)
    pipe.unload_lora_weights()
    
    pipe.save_pretrained(output_path)
    print(f"병합된 모델 저장: {output_path}")

마무리

LoRA는 파인튜닝의 황금 표준이다. 전체 파인튜닝 대비 10배 적은 자원으로 80% 수준의 품질을 달성한다. 상업적 활용 시 베이스 모델 라이선스를 확인해야 한다: SD 1.5(CreativeML OpenRAIL-M), SDXL(SDXL 라이선스). 커스텀 캐릭터 생성에는 DreamBooth, 스타일 전이에는 LoRA가 적합하다.