HiDream-O1-Image 전용 150라인 LoRA 트레이너 구현 및 x0-param 최적화

One of the First Public HiDream-O1-Image LoRAs — and How to Train Your Own

shinji shimizu2026년 5월 26일9분advanced

AI 요약

Context

HiDream-O1-Image는 VAE와 별도 Text Encoder가 없는 Pixel-level Unified Transformer(UiT) 구조를 채택한 모델임. 기존 SDXL/Flux 기반 LoRA 트레이너들은 UNet 및 VAE 의존적인 구조로 인해 해당 모델의 통합 트랜스포머 아키텍처를 처리하지 못하는 한계가 존재함.

Technical Solution

Qwen3-VL 백본 기반의 Unified Transformer 구조 분석을 통한 PEFT 적용 지점 식별
Inference 코드를 역공학하여 모델의 Output Head가 Clean Image x0를 예측하는 x0-parameterization 방식임을 도출
$z_t = (1 - \sigma)x_0 + \sigma(8.0\epsilon)$ 공식을 적용한 Noised Input 생성 로직 설계
Image Token 위치에 대해 MSE(x_pred, x0) 손실 함수를 적용한 학습 루프 구현
PATCH_SIZE 32 기준의 Raw Pixel 패치화 및 t_embedder를 통한 Timestep 주입 구조 활용
low_cpu_mem_usage=True 옵션을 통한 8B 모델 로드 시 Host-RAM OOM 문제 해결

실천 포인트

- Unified Transformer 모델 학습 시 기존 Diffusers 기반 툴셋의 호환성 여부를 먼저 검토할 것 - Flow Matching 기반 모델의 경우 Prediction Target이 v-param인지 x0-param인지 구분하여 Loss 함수를 설정할 것 - 대규모 모델 로드 시 CPU RAM 부족을 방지하기 위해 Shard 스트리밍 로딩 옵션을 적용할 것 - 입력 이미지 해상도를 모델의 Patch Size(예: 32) 배수로 설정하여 버킷팅 오차를 제거할 것

태그

#PEFT #x0-parameterization #Unified Transformer #Flow Matching #LoRA

원문 읽기