[AI] Stable Diffusion 성능 향상

Stable Diffusion 성능 향상 왜 궁금했을까❓

Stable Diffusion을 이용하여 4장의 이미지를 생성할 때, GPU T4를 기준으로 약 3분의 시간이 소요됐다. 이미지 한 장당 약 45초의 시간이 소요된 것이다. 많은 사용자 요청을 빠르게 처리하기 위해서 서버를 증설, GPU를 업그레이드, 성능 최적화이 이뤄져야 했다. 예산적인 측면에서 서버 증설과 GPU 업그레이드는 이뤄질 수 없었기에 성능 최적화에 초점을 맞춰 작업이 이뤄졌다.

 

1. 성능 최적화 방법

Stable Diffusion XL

• SDXL 공식 문서에서 Optimizations 탭에 메모리 효율을 향상시킬 수 있는 방법이 몇 가지 소개하고 있다.

1. CPU Offload를 통해 메모리 에러 해결

• 이미지 생성을 CPU가 대신 작업을 맡아 메모리 에러에 대처하는 방법인 거 같았다.
• 현재 GPU T4를 사용중이므로 메모리 부족은 발생하지 않아 해당 방법은 사용하지 않았다.

# Offload the model to the CPU with enable_model_cpu_offload() for out-of-memory errors:

- base.to("cuda")
- refiner.to("cuda")
+ base.enable_model_cpu_offload()
+ refiner.enable_model_cpu_offload()

 

2. torch.complie() 함수를 통한 속도 향상 (torch >= 2.0)

• torch의 complie() 함수를 통해 속도 향상이 이뤄지는 것인데 모델을 사전에 컴파일하여 추후 실행에서 속도를 향상시키는 방법이다.
• 자세한 이유를 알고 싶다면 아래 링크를 통해서 알 수 있다.

PyTorch 2.0

+ base.unet = torch.compile(base.unet, mode="reduce-overhead", fullgraph=True)
+ refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)

 

3. xFormers를 설치하여 메모리와 속도 향상(torch < 2.0)

• xFormers 라이브러리를 통해 이미지 추론이나 트레이닝 과정에서 속도와 메모리 효율을 최적화를 이뤄낼 수 있다고 한다.

1번 방법의 경우 메모리 효율이 좋아지긴 하지만 CPU로 추론이 이뤄져서 이미지 생성에 느릴 것이다. 그리고 3번의 경우 당일 프로젝트는 torch 2.0 이상 버전을 사용하기 떄문에 2번의 방법으로 최적화를 수행했다.

 

2. 성능 최적화

def load_sdxl():
    print("=================== Checkpoint loaded start ===================")
    base_dir = app.utils.global_vars.base_dir
    base = StableDiffusionXLPipeline.from_single_file(base_dir + "checkpoint/stable-diffusion-xl-base-1.0.safetensors", torch_dtype=torch.float16, variant="fp16").to("cuda")
    base.scheduler = DPMSolverMultistepScheduler.from_config(base.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
    print("=================== Checkpoint loaded end ===================")
    
    print("=================== lora loaded start ===================")
    base.load_lora_weights(base_dir + "lora/detail.safetensors", adapter_name="detail")
    base.load_lora_weights(base_dir + "lora/children.safetensors", adapter_name="child")
    base.load_lora_weights(base_dir + "lora/animation.safetensors", adapter_name="animate")
    base.load_lora_weights(base_dir + "lora/pixel.safetensors", adapter_name="pixel")
    # 성능 최적화
    base.unet = torch.compile(base.unet)
    print("=================== lora loaded end ===================")

• 위에서 제공한 방법과 다르게 reduce-overhead와 fullgraph 옵션을 제거했다.
• reduce-overhead를 사용할 경우 cuda를 사용할 때 오버헤드를 줄이지만 메모리가 더 쓰인다는 점이 있었다. 해당 옵션을 사용하고 SDXL을 실행해보니 GPU 메모리가 초과되는 문제점이 발생했다. 또한, 문서에서도 오버헤드를 줄이는 것이 무조건적으로 최적화에 도움되는 것은 아니라고 소개하고 있다.
• fullgraph의 경우 true를 하고 이미지 생성을 해봤더니 에러가 발생했다. 찾아보니 단일 그래프를 허용하는 옵션이라는데 LoRA를 여러 개 추가해서 발생하는 문제 같아서 해당 옵션을 삭제했다.

이러한 최적화 과정을 통해 이미지 생성 시간이 약 45초에서 35초로 감소한 결과를 가져올 수 있었다. 총 4가지 화풍의 이미지를 생성하는데 3분 걸리는 것을 2분 20초로 성능 최적화를 이뤄낼 수 있었다.