[GCP] Google Cloud의 TPU와 MaxDiffusion을 활용한 이미지 생성하기

Terry Chung Hwan Ahn - 안정환

Published in

google-cloud-apac

13 min readJul 4, 2024

안녕하세요. Terry 안정환입니다.

이전 포스팅에서 Google Cloud의 G2 VM과 L4 GPU를 활용하여 Stable Diffusion을 파인튜닝하고 이미지를 생성하는 과정을 설명드렸습니다.

이전 Medium 블로그 링크 :
[GCP] Google Cloud의 G2 VM + L4 GPU를 활용하여 Stable Diffusion을 파인튜닝하고 이미지 생성하기

이번에는 한 단계 더 나아가, Google의 TPU와 MaxDiffusion을 활용하여 Stable Diffusion 모델로 고해상도 이미지를 더욱 빠르고 효율적으로 생성하는 방법을 소개하려고 합니다. 이번 포스팅에서는 TPU v5e 모델을 사용해 볼텐데요. v5e (efficiency)는 작년 11월에 정식 출시된 5세대 TPU로, 이전 모델과 비교했을 때, 가격 대비 성능이 뛰어나다는 장점이 있습니다.

Cloud TPU v5e는 2D 슬라이스 토폴로지를 지원하여, 단일 호스트 기준으로 최대 8개의 v5e 칩을 다양하게 구성할 수 있습니다. 즉, 1x1, 2x2, 2x4 슬라이스 조합을 통해 각 슬라이스에 1개, 4개 또는 8개의 칩을 할당하여 필요에 맞는 최적의 성능을 확보할 수 있습니다.

Cloud TPU v5e에 대해서는 다음 공식 블로그와 가이드 문서를 통해서 좀더 자세하게 확인하실 수 있습니다.

Google Cloud 블로그 (영문) :
How Cloud TPU v5e accelerates large-scale AI inference | Google Cloud Blog

TPU v5e 가이드 문서 :
Google Cloud TPU v5e

TPU와 함께 소개해드릴 MaxDiffusion은 Google Cloud에서 새롭게 선보이는 고성능의 효율적인 확산 모델 (Diffusion Model)의 오픈소스 레퍼런스 구현입니다. Stable Diffusion 2부터 최신 Stable Diffusion XL Lightning까지 다양한 Stable Diffusion 모델을 최적화하여 제공하며, TPU, XLA 컴파일러, Jax 프레임워크 스택에 기반하여 구축되었습니다. 2024년 4월 Next 행사에서 공개된 MaxDiffusion은 확산 모델 분야의 새로운 기준을 제시하며, 고품질 이미지 생성의 가능성을 더욱 확장하고 있습니다.

MaxDiffusion은 HuggingFace의 Diffusers 라이브러리를 기반으로 만들어졌기 때문에, HuggingFace 모델 허브에서 제공하는 다양한 모델과 데이터셋들을 Diffusers와 동일한 방법으로 쉽게 가져와 TPU 환경에서 실행할 수 있습니다.

지금부터 MaxDiffusion과 TPU v5e의 조합을 통해 Stable Diffusion 모델의 잠재력을 이끌어내는 방법을 알아보겠습니다.

#1. Google Cloud에서 v5e TPU VM을 생성

Google Cloud Console에서 Compute Engine > TPUs 메뉴를 통해 아래와 같은 UI로 TPU를 생성합니다. MaxDiffusion GitHub의 가이드는 v4–8 TPU VM을 기준으로 작성되어 있고, 2개 이상의 TPU 칩을 필요로 합니다. 따라서 이 예제에서도 TPU v5e 칩 4개가 탑재된 v5litepod-4 TPU VM을 생성하도록 하겠습니다.

또는 다음처럼 Google Cloud Shell에서 gcloud 명령어를 통해서 TPU를 생성하는 것도 가능합니다. 아래 커맨드라인 명령에서 maxdiffusion-tpu-v5e와 my-tpu-network는 원하는 이름으로 변경이 가능합니다. accelerator-type 옵션은 v$VERSION_NUMBER-$CHIP_COUNT 형식의 문자열 값을 가지므로, 예를 들어 v5litepod-4는 v5e 4개 칩을 의미합니다.

$ gcloud compute tpus tpu-vm create maxdiffusion-tpu-v5e \
  -zone=us-central1-a \
  -accelerator-type=v5litepod-4 \
  -version=tpu-ubuntu2204-base \
  -network=my-tpu-network

프로젝트에 TPU Quota (할당량)이 없는 경우는 TPU VM 생성이 불가능합니다. Google Cloud Console에서 Quota 페이지로 이동한 후 TPU v5e 항목의 할당량을 추가로 요청해야 할 수 있습니다. “TPU v5 lite pod cores for serving per project per zone” 항목으로 검색한 후, 원하는 지역의 TPU v5e 칩 개수를 요청하시면 됩니다. 자세한 할당량 관리 내용은 아래 Google Cloud 문서 페이지를 확인해 주세요.

Google Cloud Quota 가이드 문서 :
할당량 보기 및 관리 | Cloud Quotas

#2. MaxDiffusion으로 Stable Diffusion 모델 인퍼런스 테스트

생성한 TPU VM에 SSH로 접속한 후, MaxDiffusion의 GitHub repository를 clone합니다.

$ git clone https://github.com/google/maxdiffusion

다음 명령을 실행하여 MaxDiffusion 디렉토리로 이동하고 연관된 Python 패키지를 설치합니다.

$ cd maxdiffusion
$ pip3 install jax[tpu] -f \
  https://storage.googleapis.com/jax-releases/libtpu_releases.html
$ source ~/.profile
$ pip3 install -r requirements.txt
$ pip3 install .

Stable Diffusion XL Lightning 모델은 다음과 같은 스크립트를 실행하여 이미지를 생성할 수 있습니다. 실행이 끝나면 해당 디렉토리에 “image_sdxl_*.png” 형태의 이미지 파일이 만들어지게 됩니다.

$ python -m src.maxdiffusion.generate_sdxl \
src/maxdiffusion/configs/base_xl.yml \
run_name="my_run" \
lightning_repo="ByteDance/SDXL-Lightning" \
lightning_ckpt="sdxl_lightning_4step_unet.safetensors"

이미지 생성 프롬프트와 기타 설정값들은 src/maxdiffusion.configs/base_xl.yml 파일에 선언되어 있습니다.

원하는 프롬프트를 사용해서 이미지를 생성하려면 해당 파일의 prompt 부분을 변경하여 저장한 후 스크립트를 실행하면 됩니다.

# Generation parameters
prompt: "A magical castle in the middle of a forest, artistic drawing"
negative_prompt: "purple, red"

MaxDiffusion GitHub에 공개된 인퍼런스 예제의 경우는 2개 이상의 TPU 칩을 요구합니다. 따라서 v5litepod-1 와 같이 1개의 TPU 칩이 탑재된 VM의 경우는 다음과 같은 에러가 발생합니다.

AssertionError: You must have at least two devices

다른 Stable Diffusion 버전을 사용한 이미지 생성의 경우도 MaxDiffusion GitHub 페이지에 예제가 있으니 참고하시면 좋겠습니다.

MaxDiffusion GitHub 페이지 :
https://github.com/google/maxdiffusion

#3. MaxDiffusion Inference Server 샘플 코드

MaxDiffusion 인퍼런스 서버를 간단하게 구축하고 실행하는 방법을 알아보겠습니다. FastAPI를 활용하여 요청과 응답을 처리하고, 프론트엔드 클라이언트로부터 프롬프트와 함께 이미지 생성 요청이 들어오면 MaxDiffusion을 통해 Stable Diffusion XL Lightning 모델을 기반으로 해당 프롬프트에 맞는 이미지를 생성하여 반환하는 서버입니다. 해당 서버는 TPU v5e 1개 칩을 가진 v5litepod-1 TPU VM에서도 동작합니다.

서버를 구동할 TPU VM에 접속하여 샘플 서버가 포함된 코드를 다음 GitHub 리파지토리로부터 내려받습니다.

$ git clone https://github.com/terrychahn/maxdiffusion-sample

FastAPI 패키지를 설치해 줍니다.

$ pip3 install fastapi

샘플 서버는 llm-on-gke/sdxl-tpu 프로젝트를 참고하였고, 해당 프로젝트에서 Stable Diffusion 모델을 Stable Diffusion XL Lightning로 변경하였습니다.

# SD Lightning unet
LOG.info("SDXL Lightning unet set")
unet_lightning_state_dict = load_file(hf_hub_download("ByteDance/SDXL-Lightning", "sdxl_lightning_4step_unet.safetensors"), device="cpu")
flax_unet_dict = convert_pytorch_state_dict_to_flax(unet_lightning_state_dict, pipeline.unet)
params["unet"] = flax_unet_dict

다음 명령으로 서버를 실행합니다.

$ python server/main_lightning.py

일정 시간 후, 서버가 실행 완료되면 다음과 같이 포트 8000번을 listen 하는 로그가 출력됩니다.

INFO: Started server process
INFO: Waiting for application startup.
INFO: Application startup complete.
INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)

Stable Diffusion XL on GKE GitHub 페이지 :
https://github.com/llm-on-gke/sdxl-tpu

#4. Web App 샘플 코드

앞 단계에서 실행한 MaxDiffusion 서버로 프롬프트 응답을 보내는 웹 프론트엔드 앱을 만들어 보겠습니다. Streamlit을 사용하여 간단하게 한 페이지에서 프롬프트를 입력받는 부분과 이미지 생성 요청을 보내는 버튼 UI를 구성합니다.

TPU VM이 아닌 로컬에서 샘플 웹 프론트엔드 앱이 포함된 코드를 다음 GitHub 리파지토리로 부터 내려받습니다.

$ git clone https://github.com/terrychahn/maxdiffusion-sample

앱에서 MAXDIFFUSION_SERVER_URL 환경변수 값을 참조하여 호출을 하도록 구성하였습니다. 환경변수에 서버 IP와 PORT 값을 설정한 후, streamlit run 으로 해당 코드를 실행합니다.

$ export MAXDIFFUSION_SERVER_URL="http://[SERVER_IP]:8000"
$ streamlit run src/frontend/clientapp.py

다음과 같은 화면이 브라우저에서 로딩되고 나면, 원하는 프롬프트를 입력하여 이미지 생성 요청을 MaxDiffusion 서버로 보냅니다. 최초 요청의 경우에는 JIT compilation 및 초기화 과정 때문에 시간이 더 걸리지만, 두번째 요청 이후부터는 1초 미만의 짧은 시간에 이미지가 생성되는 것을 확인할 수 있습니다.

다음은 프롬프트와 생성된 이미지 예시입니다.

“A steampunk airship soaring over a lush, bioluminescent jungle at twilight.”
(황혼녘에 무성하고 생체발광하는 정글 위를 날아오르는 스팀펑크 비행선)

“A photorealistic portrait of a person with vibrant, galaxy-swirled eyes and hair that flows like nebulae.”
(은하수처럼 소용돌이치는 눈과 성운처럼 흐르는 머리카락을 가진 사람의 사실적인 초상화)

#5. Conclusion

지금까지 Google Cloud의 TPU v5e와 MaxDiffusion을 활용하여 Stable Diffusion XL Lightning 모델로 고해상도 이미지를 생성하는 방법을 자세히 살펴보았습니다. 오픈소스 MaxDiffusion을 통해 TPU 환경에서 더욱 효율적이고 간편하게 이미지 생성 작업을 수행할 수 있음을 확인하셨기를 바라며, 이번 포스팅이 Stable Diffusion 모델 활용에 새로운 가능성을 열어주는 계기가 되었으면 합니다.

Disclaimer: 본 글의 작성자는 Google 직원이지만 Google Cloud 를 공부하는 한 개인으로서 작성된 글입니다. 본 글의 내용, 입장은 Google 을 대변하지 않으며 Google 이 해당 콘텐츠를 보장하지 않습니다.

[참고자료 모음]