하루가 멀다하고 FLUX 관련 기술들이 발전하고 있습니다. 최근들어 Local에서 FLUX LoRA 학습을 많이들 진행하는 추세인데요, 일반적으로 구동에 비해 학습은 “정밀도”를 더욱 요구하는 작업이기 때문에, 절대적인 컴퓨팅 파워가 받쳐주지 않으면 아예 진행 자체가 불가합니다. 때문에 초기에는 낮은 VRAM에서는 FLUX LoRA 학습 자체가 불가했지만, 여러 우회로가 확보되면서 현재는 상대적으로 낮은 12GB의 VRAM을 탑재한 GPU에서도 로라 학습이 가능한 상황입니다.
ComfyUI Flux Trainer
기존의 SD1.5 및 SDXL LoRA의 경우 Kohya_ss에서 큰 어려움 없이 학습이 가능했습니다. 마찬가지로 Flux 로라도 Kohya_ss 및 AIToolkit 등의 프로그램을 활용해 학습이 가능하지만, 오늘은 ComfyUI 자체에서 Flux LoRA 학습을 진행하는 방법을 살펴봅니다.
오늘 학습에 사용할 커스텀노드는 ComfyUI Flux Trainer로, kijai님이 Kohya_ss 스크립트를 파인튜닝하여, ComfyUI에서도 학습 진행이 가능하도로고 배포하고 있습니다. 먼저 학습을 진행하기 위해, ComfyUI Manager에서 ComfyUI Flux Trainer를 검색하여 설치합니다. 직접 설치를 하실 분들은 주소를 활용하시면 되고, ComfyUI에서 Custom Nodes 설치가 처음인 분들은 아래 글을 참고해주시기 바랍니다.
Repo : https://github.com/kijai/ComfyUI-FluxTrainer
Flux LoRA 학습 세팅
해당 Repo 깃허브 페이지에서 example로 이동하면, 아래와 같이 Flux LoRA Train 예제 워크플로우가 업로드 되어 있습니다.
해당 워크플로우를 다운받아 ComfyUI에서 불러옵니다. 해당 워크플로우를 실행하기 위해서는 앞서 설치한 ComfyUI Flux Trainer와 함께, 아래의 2개 커스텀 노드가 추가로 필요합니다.
KJNodes 및 rgthree 커스텀노드는 꼭 학습이 아니더라도 대부분의 워크플로우에 필수적으로 사용하는 노드이므로, 이미 ComfyUI를 사용중인 분들이라면 추가로 설치를 요구하지 않을 수 있습니다. 하지만, 최근들어 Flux 관련하여 ComfyUI 자체를 비롯해 커스텀노드 기능들이 대폭 업데이트 되고 있으므로, 진행 중에 오류가 발생하는 경우, 업데이트를 진행하여 다시 시도해보시기 바랍니다.
Workflow Breakdown
Missing nodes 설치가 완료되었다면, 위와 같이 가로로 길게 짜여진 워크플로우를 마주할 수 있습니다. 각 부위별 역할을 한 번 살펴보겠습니다.
Dataset
가장 좌측에 Dataset 그룹이 있습니다. 여기서는 말 그대로 학습에 사용될 데이터를 로딩해주는 역할을 합니다. Flux 모델을 512, 768, 1024 해상도를 모두 섞어서 학습시킬 수 있습니다. 하지만, 만약 12GB VRAM으로 학습을 진행한다면, 512 해상도로 먼저 시도해보시는 것을 권장드립니다.
TrainDatasetAdd 노드 3개 중, 2개를 bypass해주신 후, 아래와 같이 512해상도만 남겨줍니다.
TrainDatsetAdd 노드의 첫번째 공란에는 학습할 이미지를 담아둔 폴더의 경로를 입력하고, 그 아래는 클래스 토큰을 작성합니다. 따로 떠오르는 단어가 없다면 공란으로 두셔도 됩니다. enable_bucket을 true로 하셨다면, 준비한 이미지의 크기가 각기 다르더라도 알아서 해상도에 맞게 학습을 진행합니다. num_repeats는 학습 진행시 적용되는 리피트 값을 의미합니다.
Settings and Init
본격적으로 학습을 위한 설정을 해주는 그룹입니다. SD1.5 또는 SDXL에서 학습을 진행할 때 다뤘던 익숙한 설정이므로, 아래 몇가지 주의사항만 참고하여 적당히 본인의 취향에 맞게 설정하시면 됩니다.
FluxTrain ModelSelect에서는 기본 설정된 모델들을 동일하게 세팅해주시면 됩니다. Flux dev 본모델(FP16)을 설정하더라도, VRAM 12GB에서는 학습을 진행할 때 FP8 모드로 진행이 되는 것을 확인했습니다. 따라서 VRAM 12GB에서 학습 진행시에는 flux dev fp8, flux vae, clip_l, t5xxl 모델을 설정해줍니다. Flux Train Validation Settings는 따로 건들지 않아도 됩니다.
2가지 주의사항 중 첫번째는 Optimizer Config를 Adafactor로 바꿔 연결하는 것과, split_mode를 true로 설정해주시는 것입니다. VRAM 12GB GPU에서 Flux LoRA를 학습하기 위한 설정입니다. 스플릿 모드를 통해 부족한 VRAM을 극복하고, 학습을 진행하게 됩니다. 그외 설정은 일반적인 LoRA 학습 설정과 동일하게 진행합니다. network_dim과 network_alpha 값은 보통 동일한 값으로 맞추며, 32, 64, 128, 256 등의 값을 입력하는데요, dim 값이 높을수록 생소하고 복잡한 개념을 학습하는 효율이 올라갑니다. 반면 그만큼 LoRA의 크기도 증가합니다.
Max_train_steps를 통해 총 몇스텝을 진행할지 결정할 수 있습니다. 기본 설정으로 3,000이 설정되어 있으며, 현재 설정을 기준으로, 인물 학습의 경우 보통 1,500 ~2,000 스텝 사이에서 적정 수준의 학습이 이뤄진다고 알려져 있습니다.
Train
이제 마지막으로, 뒷단을 따르는 4개의 Train 그룹에서 Flux Train Loop 노드의 steps를 설정해주면 됩니다. 여기서 설정한 스텝만큼의 학습이 진행되면, 해당 스텝에서 LoRA를 저장하게 됩니다. 즉, 750스텝씩 네번의 Train 그룹을 거쳐 총 3,000 스텝(=750 x 4)의 학습을 완료하게 되며, 이때 4개의 로라 파일(750, 1500, 2250, 3000)이 생성됩니다. 상황에 맞게 적절한 값을 입력하여 학습을 진행할 수 있습니다.
이제 모든 설정이 완료되었으므로, Queue 프롬프트를 눌러 ComfyUI에서 Flux LoRA 학습을 진행할 수 있습니다. 참고로 RTX3060 12GB를 기준으로 원본이미지 18장, 600스텝 학습을 진행하는데 약 2시간 반정도가 소요되었으며, 600스텝 로라를 사용했을 때, 꽤 그럴듯하게 원본과 유사한 인물을 생성해낼 수 있었습니다.
