[Train] ControlNet 학습시키기

 

https://github.com/lllyasviel/ControlNet/blob/main/docs/train.md 

ControlNet/docs/train.md at main · lllyasviel/ControlNet

 

ControlNet을 학습시켜보자. 이 문서에 잘 나와있긴하지만, 몇가지 디버깅했던 내용을 포함해서 적어보려한다. 

ControlNet은 github, huggingface 모두 다 올라와있는데, github에는 코드 중심으로, huggingface에는 데이터 중심으로 나와있어서 서로 다르다.

 

하지만 모든 코드는 git이 base이기때문에 git clone을 하고 시작하는 것을 추천한다.

git clone https://github.com/lllyasviel/ControlNet.git

 

1. Dataset Download하기

데이터셋의 경우 ControlNet Huggingface에 zip 파일이 올라와있다. 

terminal 상에서 

 

wget "https://huggingface.co/lllyasviel/ControlNet/resolve/main/training/fill50k.zip"

을 하여 다운받을 수 있다. 

그리고 zip파일이므로 압축을 풀어준다.

 

unzip training/fill50k

그러면 다음과 같이 압축이 풀린 데이터들이 나온다. 

 

2. Dataset Load하기

이제 다운받은 데이터셋을 load해야한다.

 

https://github.com/lllyasviel/ControlNet/blob/main/tutorial_dataset.py

ControlNet/tutorial_dataset.py at main · lllyasviel/ControlNet

이 tutorial_dataset.py는 앞서 다운받은 fill50k 폴더 안에 있는 prompt.json파일에서 데이터 정보를 읽어와서, 

각 아이템을 jpg(타겟이미지), txt(프롬프트), hint(소스이미지)를 포함하는 딕셔너리 형태로 반환하는 역할을 한다. 

ControlNet은 소스이미지를 입력으로 사용하고 텍스트 프롬프트(txt)를 조건으로 사용하여 타겟이미지(jpg)를 생성하는 것을 목표로 삼아 학습한다. 여기서 왜 타겟이미지가 jpg라고 불리고, 소스이미지가 hint라고 불리는지는 그냥 convention인듯싶다. 아무이유없는 것 같다. (왜냐하면 이 데이터셋에서는 이미지가 png이다.)

 

아무튼

 

이 tutorial_dataset.py가 정의된 상황에서 

https://github.com/lllyasviel/ControlNet/blob/main/tutorial_dataset_test.py

ControlNet/tutorial_dataset_test.py at main · lllyasviel/ControlNet

tutorial_dataset_test.py 코드를 다음과 같이 실행시키면 

 

python tutorial_dataset_test.py

 

아래와 같은 output을 확인할 수 있다. 

 

50000
burly wood circle with orange background
(512, 512, 3)
(512, 512, 3)

 

이 결과가 나왔으면 dataset은 잘 로드가 되는 상태가 되었네, 하고 넘어가면 된다.

 

3. ControlNet을 어떤 StableDiffusion 모델에 연결할지 결정

다들 잘 알다싶이, ControlNet은 Stable Diffusion 모델과 연결하고, 그 Stable Diffusion model의 parameter를 copy한 상태에서 학습을 시작한다. Stable Diffusion Model이 버전이 여러개가 있으니, 어떤 것을 사용할지 정해야한다. 

이 Stable Diffusion Model의 가중치는 아래의 HuggingFace에서 가져올 수 있다. 

 

https://huggingface.co/runwayml/stable-diffusion-v1-5/tree/main

runwayml/stable-diffusion-v1-5 at main

 

 

나의 경우 v1-5-pruned.ckpt를 가지고 왔다. 

 

wget "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned.ckpt"

 

꽤 용량이 크다. 내 경우 7.7GB가 필요했다. 아무튼, 이 가중치 파일을 models폴더에 넣어놓고 ControlNet을 Stabld Diffusion model

에 연결하는 스크립트를 실행해야한다. 

 

 

tool_add_control.py 는 https://github.com/lllyasviel/ControlNet/blob/main/tool_add_control.py 에서 찾을 수 있다. 

 

스크립트 내용은

'cldm_v15.yaml' 설정 파일을 사용하여 ControlNet 모델을 생성하고, 입력 경로에서 사전 학습된 가중치를 로드하고,

새 모델의 구조에 맞게 사전 학습된 가중치를 매핑한다.
'control_'로 시작하는 레이어는 'model.diffusion_'으로 시작하는 가중치와 매핑된다. 매핑되지 않은 새 가중치는 초기화된 상태로 유지되며 출력된다. 매핑된 가중치를 새 모델에 적용한다. 최종 모델의 상태를 지정된 출력 경로에 저장하는 스크립트이다. 

 

 

 

python tool_add_control.py ./models/v1-5-pruned.ckpt ./models/control_sd15_ini.ckpt

 

이를 실행하면, 다음과 같은 에러가 뜰 것이다. 

 

  File "/local_data_2/home/ControlNet/ldm/models/diffusion/ddpm.py", line 20, in <module>
    from pytorch_lightning.utilities.distributed import rank_zero_only
ModuleNotFoundError: No module named 'pytorch_lightning.utilities.distributed

 

 PyTorch Lightning의 특정 모듈을 찾을 수 없다는 말인데, PyTorch Lightning의 버전 차이때문에 나타나는 오류일듯하다. 

이는 https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/11458

[Bug]: ModuleNotFoundError: No module named 'pytorch_lightning.utilities.distributed' · Issue #11458 · AUTOMATIC1111/stable-di

이 이슈의 도움을 받아 해결할 수 있는데, 간략하게 말하면

'ldm/models/diffusion/ddpm.py' 파일에서 import를 해오는 부분인 20번째 줄에서 'pytorch_lightning.utilities.distributed'를 'pytorch_lightning.utilities.rank_zero'로 변경하면 해결된다. 아래와 같이 20번째 줄이 적혀있으면 그 에러는 생기지 않는다. 

 

그러면 다시, 앞서 설명했던

python tool_add_control.py ./models/v1-5-pruned.ckpt ./models/control_sd15_ini.ckpt

 

를 실행시키면 되겠다. 정상적으로 실행이 된다.

ControlNet 모델을 초기화하고 가중치를 로드하는 과정을 볼 수 있다. ControlNet의 새로운 구성 요소들로 기존

모델에 없던 새로운 가중치 레이어가 추가되는 모습을 볼 수 있다. 이 추가된 가중치들은 ControlNet의 구성요소인 Control model의 일부이다. 

Zero Convolution의 weight와 bias, 그리고 input_hint_block과 middle_block_out의 가중치 레이어가 추가되었음을 알 수 있다. 

 

1) 처음에 추가된 logvar은 variational inference에 사용되는 로그 분산 파라미터로 모델의 불확실성을 조절하는데 사용될 수 있다. 

 

2) Zero Convolution은 처음부터 너무 control에 해당하는 부분이 급격하게 학습되는 것을 막기위한 합성곱 레이어이다. 

 

3) Input_hint_block은 제어입력을 처리하는 초기 블록이다. 여러개의 합성곱 레이어로 구성되어있으며 입력된 hint를 모델이 사용할 수 있는 형태로 변환한다.

 

4) Middle_block_out은 ControlNet의 중간 처리 단계에서 사용되는 출력 블록으로 처리된 제어신호를 Stable Diffusion의 중간 레이어와 통합하는 역할을 한다. 

 

이를 마치면 앞서 작성했던 ./models/control_sd15_ini.ckpt 으로

Stable Diffusion과 ControlNet이 통합된 전체 모델의 가중치가 포함된 체크포인트 파일이 저장된다. 

 

4. 학습시키기

이제 학습시킬 수 있는 상태가 되었다. 

https://github.com/lllyasviel/ControlNet/blob/main/tutorial_train.py

ControlNet/tutorial_train.py at main · lllyasviel/ControlNet

 

이 파일을 실행시키면 된다. 사실 이 파일을 실행시키려면 cldm.model 이나 cldm.logger 파일등이 필요하겠지만 맨 처음에 git clone했을때 잘 복사해왔을거라 가정하고 진행하겠다. 

 

python tutorial_train.py

 

이렇게 훈련을 시키면 또 다음과 같은 에러가 뜰 것이다. 

 

ModuleNotFoundError: No module named 'pytorch_lightning.utilities.distributed'

 

그러면 또 저 에러가 난 파일에 들어가서 (내 경우, cldm/logger.py", line 8 였음.)

수정 후 파일

'pytorch_lightning.utilities.distributed'를 'pytorch_lightning.utilities.rank_zero'로 변경하면 해결이 된다. 

 

그런데 또 에러가 난다.

 

TypeError: LatentDiffusion.on_train_batch_start() missing 1 required positional argument: 'dataloader_idx'

 

이 오류는 PyTorch Lightning의 최신 버전과 LatentDiffusion 클래스의 on_train_batch_start 메서드 정의 사이의 불일치로 인해 발생한듯 한데, 

https://github.com/lllyasviel/ControlNet/issues/84

colab error, pytorch lighting trainer · Issue #84 · lllyasviel/ControlNet

이 이슈를 참고하면 문제가 해결된다. 이또한 간략하게 말하자면 ldm/models/diffusion/ddpm.py 의 591번째 on_train_batch_start 함수 정의부분에 가서 dataloader_idx부분을 삭제해주고, cldm/logger.py의 74번째 on_train_batch_end 함수 정의 부분에 가서도 dataloader_idx 부분을 삭제해주면 에러는 해결될 것이다.

 

하지만 마지막 에러가 남았다.

 

TypeError: Trainer.__init__() got an unexpected keyword argument 'gpus'

 

코드가 나온 직후에 많은 업데이터가 되었겠거니.. 라는 마음을 가지고 

tutorial_train.py의 29번째 줄을 다음과 같이 변경해준다. 

 

Before 

trainer = pl.Trainer(gpus=1, precision=32, callbacks=[logger])

 

After

trainer = pl.Trainer(accelerator="gpu", devices=1,precision=32, callbacks=[logger])

 

이렇게 바꿔서 실행하면, 학습이 진행된다.

총 1.4B 개의 파라미터 중 1.2B 개가 학습 가능한 것 같다. 

 

 

+) 마지막 에러
ImportError: cannot import name 'get_full_repo_name' from 'huggingface_hub'
이 에러는 

 

https://github.com/comfyanonymous/ComfyUI/issues/2055

ImportError: cannot import name 'get_full_repo_name' from 'huggingface_hub' · Issue #2055 · comfyanonymous/ComfyUI

 

이 깃허브 이슈를 참고해서 conda install chardet 로 해결해줄 수 있었다.