[Generative] Diffusion 모델 원하는 위치에 받기

<분류>
A. 수단
- OS/Platform/Tool : Linux, Kubernetes(k8s), Docker, AWS
- Package Manager : node.js, yarn, brew, 
- Compiler/Transpillar : React, Nvcc, gcc/g++, Babel, Flutter
- Module Bundler  : React, Webpack, Parcel

B. 언어
- C/C++, python, Javacsript, Typescript, Go-Lang, CUDA, Dart, HTML/CSS

C. 라이브러리 및 프레임워크 및 SDK
- OpenCV, OpenCL, FastAPI, PyTorch, Tensorflow, Nsight

 

---

1. What? (현상)

 

Diffusion을 활용해 생성 모델을 구축하는데 가장 많이 활용하는 패키지는 Diffusers일 것입니다. 

** https://huggingface.co/docs/diffusers/index

 

Diffusers는 Huggingface에서 만든 패키지로 이미지, 오디오, 3D 정보 등을 pre-trained 모델을 활용해 생성하기도 하고, 학습을 하기위해서도 활용합니다. 

 

이때 Diffusion 모델들을 fine-tuning을 할 때도 학습이 오래걸리기 때문에 pre-trained 모델을 활용하고, Inference를 할때도 역시 학습된 모델들을 활용해야 합니다.

 

이때 이 모델을 구하는 곳이 바로 Huggingface(https://huggingface.co/)와 Citivtai(https://civitai.com/)입니다.

 

학습에 사용하는 경우, 직접 repository에서 받아 학습이 가능하도록 구현할 수 있지만, ComfyUI나 Kohya와 같이 직접 UI를 통해 모델을 제공해야하는 경우 로컬에 직접 저장하는 방법도 알아두면 좋을 것 같습니다.

 

이번 글에서는 pre-trained된 모델들을 로컬에 원하는 위치에 다운로드하는 방법에 대해 정리하고자 합니다.

---

2. Why? (원인)

• X

---

3. How? (해결책)

 

0. Diffusers에서 학습하는 예시

 

Diffusers를 활용해 학습을 활용할 때 보통 특정한 구조의 폴더 Tree를 필요로 합니다. 이를 먼저 살펴보고자 합니다.

 

아래는 Stable Diffusion 모델을 불러온 경우의 폴더트리의 예시입니다. 아래 예시는 공식 Stable Diffusion 모델(https://huggingface.co/runwayml/stable-diffusion-v1-5)을 보였습니다.

• feature_extractor : (Safety-check) 생성된 이미지를 safety_checker를 통해 확인하기 위해 feature를 얻어내는 모델입니다.
• safety_checker : (Safety-check) Classification 모델을 활용해 생성된 이미지가 offensive하거나 harmful한지를 찾아내는 SD의 모듈입니다.
• scheduler : (Diffusion) encoded된 image latents를 denoise할 때 UNet과 함께 사용되는 Scheduler입니다. DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler 등을 활용할 수 있습니다.
• unet : (Diffusion) encoded된 image latents를 denoise할 때 사용하는 UNet입니다.
• text_encoder : (Control) Frozen된 text-encoder로, SD에서는 CLIP을 활용해 합니다.
• tokenizer : (Control) 위 text_encoder에서 활용하는 tokenizer입니다.
• vae : (Encoder) image를 latent representation으로 encode/decode하기 위한 Variational Auto-Encoder(VAE) 모델입니다.
• model_index.json : pipeline클래스를 통해 불러올 때, 어떤 클래스를 활용하며 어떤 모듈들을 불러올 것인지에 대한 정보를 담은 config파일입니다.

위와 같은 폴더트리를 가진 Stable Diffusion 모델은 아래와 같은 명령어로 불러올 수 있습니다. 아래는 로컬에 stable-diffusion-v1-5라는 폴더가 이미 구해져있는 경우의 예시를 보였습니다.  DiffusionPipeline 클래스는 model_index.json 파일을 통해 체크포인트에 대응되는 적합한 파이프라인 클래스와 컴포넌트들을 불러옵니다.

** .safetensors : 일반적으로 PyTorch 모델 weight는 python에서 제공하는 pickle을 활용해 .bin 파일에 저장됩니다. 하지만 이런 파일에는 악성 코드 등이 포함될 수 있는 위험이 있어 대안으로 활용되는 huggingface의 파일 포맷입니다.

from diffusers import DiffusionPipeline

repo_id = "./stable-diffusion-v1-5"
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id, use_safetensors=True)

 

혹은 아래와 같이 각각의 모듈을 직접 불러와 사용할 수도 있습니다. 아래는 Huggingface의 Repository에서 불러오는 경우인데, 불러온 모델은 ~/.cache/huggingface/hub 에 저장되었을 것입니다.

from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_safetensors=True)
tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(
    "CompVis/stable-diffusion-v1-4", subfolder="text_encoder", use_safetensors=True
)
unet = UNet2DConditionModel.from_pretrained(
    "CompVis/stable-diffusion-v1-4", subfolder="unet", use_safetensors=True
)

 

---

 

그럼 이제 이 모델들을 내가 원하는 위치에 저장하는 방법들을 살펴보겠습니다.

 

1. Diffusers 활용하기

 

위에서 설명한 바와 같이 Diffusers 패키지에 repository를 명시해주면 모델을 로컬에 알아서 받아 줍니다. 따라서, 다운받은 파일을 원하는 위치에 다시 저장해주기만 하면 됩니다.

 

위에서 예시로 든 Stable Diffusion모델(https://huggingface.co/runwayml/stable-diffusion-v1-5)을 가져와 저장해보겠습니다. 위에서 사용한 from_pretrained()함수를 활용하면 모델을 다운로드 받을 수 있습니다.

** StableDiffusionPipeline Class : DiffusionPipeline Class의 하위 모듈로 Stable Diffusion을 불러오기에 적합한 클래스 입니다.

 

그다음 save_pretrained()함수를 통해 특정 폴더에 해당 모델을 저장할 수 있습니다.

from diffusers import StableDiffusionPipeline
import torch
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.save_pretrained("./folder_name")

----------------------------------------------------------------

<이외 필자가 다운로드 받은 파일의 예시>

 

1. https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0

 

 Stable Diffusion의 XL모델로, 가장 많이 사용하는 Repostiory입니다.

아래와 같은 명령어를 통해 받았습니다.

from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")

** variant : 다른 버전의 모델 weight를 받기 위한 argument로, 예를 들어 diffusion_pytorch_model.bin라는 모델이 있을 때, variant="fp16"이라는 argument를 넘겨주면 diffusion_pytorch_model.fp16.bin라는 weight를 받게 됩니다.

 

 2.https://huggingface.co/lllyasviel/sd_control_collection

 

ControlNet으로 사용하기 위한 모델을 모아놓은 Repository입니다.

ControlNet 뿐 아니라 t2i-adapter_xl_sketch와 같이 T2I-Adapater등을 위한 모델도 가지고 있습니다.

 

----------------------------------------------------------------

 

---

2. git을 활용해 받기

 

아래와 같이 git 을 활용해 직접 repository에서 받을 수도 있습니다. 대용량 다운로드가 필요할 수도 있으므로 Git LFS(Larage File Storage)를 미리 설치해준 후 clone 해 줍니다.

apt-get install git-lfs # or git lfs install
git clone https://huggingface.co/runwayml/stable-diffusion-v1-5

 

 

---

3. Civitai에서 가져오기

위 방법외에도 실제 Diffusion모델을 모아놓은 오픈 소스 웹사이트가 많습니다. 그 중 Civitai라는 페이지에서 가져오는 것을 예로 들어보겠습니다.

 

Realistic Vision(https://medium.com/@natsunoyuki/using-civitai-models-with-diffusers-package-45e0c475a67e)이라는 모델을 다운로드받는 상황을 가정해보겠습니다.

** Realistic Vision : Stable Diffusion기반으로 학습되어 실사 이미지를 생성하는 데 있어 가장 퀄리티가 괜찮은 모델입니다.

 

아래와 같이 내가 원하는 모델의 다운로드링크의 주소를 복사합니다.

 

그 다음 아래와 같은 명령어를 통해 모델을 다운로드 받습니다.

** --content-disposition : HTTP 응답에서 활용하는 Content-Disposition 헤더와 같은 뜻으로, 컨텐츠가 브라우저에 inline 되어야 하는 웹페이지 형태인지, 아니면 다운로드 해 로컬에 저장할 attachment인지를 알려주는 헤더입니다.

wget https://civitai.com/api/download/models/130803 —content-disposition

 

이렇게 다운로드 받은 파일은 아마 .safetensors하나의 파일로 제공될 확률이 높습니다. 그 모델이 SD기반 혹은 SD-XL기반의 모델인 경우 Diffusers에서 활용하기 위해 포맷을 바꾸어 주어야하는데, 아래와 같은 명령어로 포맷을 바꾸어 주면 활용할 수 있습니다.

python3 convert_original_stable_diffusion_to_diffusers.py 
	--checkpoint_path downloaded.safetensors 
	--dump_path after_folder/ 
	--from_safetensors

 

convert_original_stable_diffusion_to_diffusers.py 파일을 구하는 방법에 대해서는 아래 더보기를 참조하세요. 

---------------------------------------------------------------------

 <Git에서 convert_original_stable_diffusion_to_diffusers.py 파일 가져오는 방법>

 

위 소스를 가지고 있는 git page(https://github.com/huggingface/diffusers/blob/main/scripts/convert_original_stable_diffusion_to_diffusers.py)에 들어간 뒤 아래 그림과 같이 Raw버튼을 클릭하면 해당 소스코드를 확인할 수 있습니다.

[소스코드]

 

Raw버튼을 누른 후에 나오는 주소를 복사해 아래와 같이 가져올 수 있습니다.

wget https://raw.githubusercontent.com/huggingface/diffusers/main/scripts/convert_original_stable_diffusion_to_diffusers.py

 

추가로 위와 반대로 convert_diffusers_to_original_stable_diffusion.py라는 파일을 아래와 같이 활용해 폴더를 하나의 safetensors로 만들어 줄 수도 있습니다.

python3 convert_diffusers_to_original_stable_diffusion.py 
	--model_path=before_folder 
	--checkpoint_path=after
	--use_safetensors

 

---------------------------------------------------------------------

 

---

 

4. Huggingface를 활용

 

huggingface 자체에서도 다운로드 할 수 있는 방법을 제공합니다.

 

먼저 snapshot_download 함수는 전체 repository를 특정해 다운로드 받을 때 활용합니다. 이때 cache_dir를 원하는 위치에 지정해 다운로드 받을 수 있습니다. 또한 revision을 통해 버전을 repository의 버전을 명시해줄 수도 있습니다.

from huggingface_hub import snapshot_download
snapshot_download(repo_id="lysandre/arxiv-nlp", cache_dir="./folder_name", revision="main")

 

다음으로 hf_hub_download 함수는 하나의 파일을 특정해 다운로드 받을 때 활용합니다. repository를 제공하고 filename을 명시해서 다운로드 합니다. 역시나 cache_dir를 원하는 위치에 지정해 다운로드 받을 수 있습니다. 

from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="stabilityai/stable-diffusion-xl-base-1.0", filename="sd_xl_base_1.0.safetensors", revision="main", cache_dir="./folder_name")

 

---

https://huggingface.co/docs/huggingface_hub/guides/download