Stable diffusion을 업무에 활용하기 #4 캐릭터 전용 학습 툴 제작기

얼마 전 Kohya를 이용해 학습하는 메뉴얼을 사내에 뿌려보았으나 아무도 안씀. -_-a 이전 글을 보시거나 이미지 학습을 수행해 보신 분들은 아시겠지만 알아야 할 설정값이 너무 많았습니다. 툴 설치나 관리도 문제고...

이대로는 아무도 안쓸 것 같아 (#4에서는 colab으로 학습하는 방법을 로컬에서 사용할 수 있게 변경해서 다시 글을 올려보겠습니다. 라고 이미 계획은 있었지만 ) 그냥 실행파일 하나 돌리면 LoRA 파일이 생성되도록 다시 작업을 했습니다.

소스코드만 필요하신 분들은 깃헙에서 바로 코드를 확인하실 수 있습니다.

---

학습에 필요한 각종 파라미터의 상세한 설명은 #3에서 충분히 설명했으니 이곳에서는 windows os에 nvidia 그래픽 카드를 가진 시스템에서 손쉽게 학습할 수 있는 시스템을 구축하는데 중점을 두고 설명합니다. 

Python 혹은 Colab을 잘 다루지 못하는 동료를 위해 로컬 윈도우 OS PC에서 손쉽게 LoRA를 생성할 수 있도록 Kohya-colab 코드를 수정했습니다. 현재는 SD 버전에서만 테스트 되었습니다.

Kohya-colab의 코드와 거의 유사하며 windows os에서 단독으로 실행하기 위해 일부 코드만 수정되었습니다. 여기서는 수정된 코드와 로컬에서의 실행을 위한 가이드만 기술하니 보다 상세한 내용은 위 링크를 참고하세요.

사전 설치할 프로그램들

 : python 3.10.x
 : sed - https://gnuwin32.sourceforge.net/packages/sed.htm
 : cuda - https://developer.nvidia.com/cuda-toolkit
 : cudnn - https://developer.nvidia.com/cudnn
 : aria2c - https://aria2.github.io/
 : 코드 설치   : git clone https://github.com/blackwitch/kohya-colab2local.git
 : 패키지 설치 : pip install -r requirements.txt

실행 절차

- config 설정 : 
  resolution - 학습할 이미지 크기값
  cuda_home - 설치된 cuda 경로 설정
  lora_name - lora 파일 이름
  model_url - 설정, 학습할 모델 주소 설정-기본값은 https://huggingface.co/hollowstrawberry/stable-diffusion-guide/resolve/main/models/animefull-final-pruned-fp16.safetensors, civit의 모델을 지정하려면 https://civitai.com/models/7240 과 같이 입력하면 됨. 
  activation_tag -  프롬프트에 추가할 lora 전용 테그값. 
  remove_tags -  자동으로 생성될 프롬프트에서 삭제할 텍스트를 설정.
- dataset 준비: 
  python run_dataset.py (아직 학습할 이미지를 추가하기 전에 기본 폴더 구조와 config 파일을 생성합니다.)
  "Loras\datasets\your_lora_name"에 학습할 이미지를 넣으세요.
  python run_dataset.py
  datasets 폴더에 생성된 txt 파일을 열어 프롬프트를 수정할 수 있습니다.
- training :
  python run_training.py
  (* triton lib 에러는 무시하세요. linux lib 입니다.)

수정된 코드 설명

 간편한 실행을 위해 학습을 위한 파라미터들은 모두 기본값으로 설정되어 있습니다. 필요한 경우 코드 내 값을 수정하여 사용하세요. 

아래에는 윈도우 환경에서 간단히 실행하기 위해 수정된 주요 내용을 정리해 보았습니다. 실행 시 자주 수정이 필요한 몇 가지 설정을 config.toml에서 설정할 수 있도록 했습니다. 위에 설명된 "config 설정"을 참고하세요.

 colab을 위해 필요한 IPython, IPython.display를 사용하지 않으므로 import하지 않습니다. 또한 !으로 실행되는 명령은 윈도우를 위해 subprocess.run으로 대체되었습니다. 이와 관련된 내용은 아래 설명에서 중복으로 설명하지 않겠습니다.
 
[run_dateset.py 파일]
 (source) https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Dataset_Maker.ipynb 
 
- "STEP 1: Setup"에서는 Google Drive를 대신해 로컬 드라이브를 사용할 수 있도록 관련 코드가 일부 추가되었습니다.

if getattr(sys, 'frozen', False):
    # frozen 속성이 있으면 실행 파일이므로 sys.executable의 디렉토리 경로를 가져옴
    current_directory = os.path.dirname(sys.executable)
else:
    # frozen 속성이 없으면 스크립트가 실행 중이므로 현재 파일(__file__)의 디렉토리 경로를 가져옴
    current_directory = os.path.dirname(__file__)

data_path = os.path.join(current_directory, "Loras")

- "STEP 2: Scrape images from Gelbooru" 은 사용하지 않아 주석처리 했습니다. 
- "STEP 3: Curate your images" 또한 이미 선별된 이미지를 사용해 학습을 진행하므로 주석처리 했습니다.
- "STEP 4: Tag your images" 에서는 로컬 실행 환경에서는 불필요한 경로 변경 및 설치 코드를 삭제했으며 그 외에는 큰 수정 사항이 없습니다.
- "STEP 5: Curate your tags"는 변경 사항이 없습니다.

[run_trainer.py 파일]
 (source) https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb

- "STEP 1: Setup"에 cuda 사용할 수 없을 경우 실행되지 않게 코드가 추가되었습니다. 

os.environ["CUDA_HOME"] = main_config.get("general").get("cuda_home")

print("cuda version =", torch.version.cuda , "and available =", torch.cuda.is_available())
if not torch.cuda.is_available():
    print("이 시스템에서는 CUDA를 사용할 수 없습니다. 재설치 후 다시 실행하세요.")
    exit(1)

고정된 값으로 사용할 툴이라 model_url 설정을 간속하게 설정했습니다. optional_custom_training_model_url을 이용할 경우 실행할때마다 다운받을 수 있으니 반복적으로 사용하려면 optional_custom_training_model_url값은 비워두고 코드에서 model_url값을 직접 설정하세요.

optional_custom_training_model_url = main_config.get("general").get("model_url")
.
.
if optional_custom_training_model_url:
    model_url = optional_custom_training_model_url
else:
    model_url = "https://huggingface.co/hollowstrawberry/stable-diffusion-guide/resolve/main/models/animefull-final-pruned-fp16.safetensors"

- "STEP Processing, Steps, Learning, Structure"는 코드에 변화는 없습니다. 옵션값의 변화가 필요하면 여기 코드를 직접 수정하세요. 가능하면 최적의 값을 찾아 세팅 후 수정을 가하지 않는 것이 사용자를 위해 좋습니다.
- 코드 하단의 경로 설정, clone_repo 등은 로컬 경로 사용을 위해 코드 일부가 수정되어 있습니다. 윈도우 환경에서 bitsandbites 사용을 위해 dll 및 소스 코드를 설정하는 부분은 아래와 같이 수정되어 있습니다.

# .venv 폴더가 있는지 확인
venv_folder = os.path.join(current_directory, ".venv")
if os.path.exists(venv_folder) and os.path.isdir(venv_folder):
	target_folder = "..\\..\\.venv\\Lib\\site-packages\\bitsandbytes\\"
else:
	target_folder = sys.prefix + "\\Lib\\site-packages\\bitsandbytes\\"

subprocess.run(["copy", ".\\bitsandbytes_windows\\*.dll", target_folder], shell=True)
subprocess.run(["copy", ".\\bitsandbytes_windows\\cextension.py", f"{target_folder}cextension.py"], shell=True)
subprocess.run(["copy", ".\\bitsandbytes_windows\\main.py", f"{target_folder}cuda_setup\\main.py"], shell=True)

 

여기까지가 코드 수정에 대한 설명입니다. 핵심 코드의 수정은 거의 없으며 windows os 환경에 맞게 수정된 부분이 대부분이니 기존 코드를 보신 분이면 이해하는데 큰 어려움은 없을 듯 합니다. 

사용에 어려움이 있으신 분들은 코멘트 남겨주세요. :)