요약하자면, 이 블로그에서는 인기 있는 3D 모델인 TRELLIS.2를 사용자의 3D 에셋에 맞춰 파인튜닝할 수 있는 스크립트와 컨테이너를 제공합니다. 이 시스템은 로컬 환경과 Hugging Face Jobs 모두에서 실행 가능합니다. 만약 파인튜닝에 관심이 있다면, 몇 가지 인자만 사용하여(데이터셋을 Hugging Face Hub로 업데이트했다고 가정) 저희의 스크립트를 사용하거나 이 저장소에서 직접 수정해 볼 수 있습니다. 더 자세히 알고 싶다면, 스크립트가 어떻게 작동하는지 살펴보겠습니다. 시작합시다!
3D 생성은 로보틱스, 시뮬레이션, 미디어 및 엔터테인먼트 등 다양한 분야에 활용됩니다. 이러한 분야에서는 생성된 객체가 엄격한 규칙을 따르고 모델과 데이터 표현이 표준화되기 어렵습니다. GLB나 PLY와 같은 수많은 파일 형식이 존재하며, 신경 방사 필드(NeRFs)나 3D Gaussian splatting (3DGS)과 같은 다양한 표현 방식도 있습니다. 각각은 서로 다른 장점과 단점을 가지고 있습니다.
TRELLIS는 통합 잠재 표현(unified latent representation)이라는 공통 내부 표현을 도입했습니다. 이 디코더들은 해당 표현을 필요한 출력 형식의 3D 모델로 변환합니다. 따라서 각 표현 형식마다 다른 모델을 사용할 필요가 없습니다.
같은 팀은 이후 TRELLIS.2를 출시했는데, 이는 모델이 3D 구조를 표현하는 능력을 향상시켰습니다. 하지만 두 모델 모두 특정 작업에 맞게 적응할 수 있는 쉬운 방법을 제공하지는 않습니다.
TRELLIS.2와 같은 범용 모델은 특정 애플리케이션에 중요한 세부 사항에서 오류를 일으킬 수 있습니다. 예를 들어, 사용자는 가구 모델의 정확한 기하학적 구조가 필요할 수 있습니다. 해당 작업에 대한 파인튜닝 없이는 사전 학습된 모델이 충분히 정확하지 않은 기하학적 구조를 생성할 수 있습니다. 저희는 실험을 위해 188개의 신발 관련 객체가 포함된 GSO의 하위 집합을 사용했습니다. 초기 테스트에서, 사전 학습된 모델은 보통 올바르게 보이는 신발을 생성했지만 밑창(sole) 부분은 정확하지 않았습니다. 우리는 LoRA를 사용하여 밑창 부분을 더 정확하게 만들 수 있는지 테스트합니다.
우리의 목표는 Google Scanned Objects (GSO) 데이터셋을 사용하여 특정 객체 카테고리에서 TRELLIS.2의 성능을 향상시키는 것입니다. 우리는 훈련 중에 어떤 파이프라인 구성 요소를 미세 조정하고 어떤 것을 변경하지 않을지 설명할 것입니다. 이 파이프라인은 구조(structure), 상세 기하학(detailed geometry), 재료(materials)를 분리합니다. 따라서 우리는 우리의 작업에 중요한 구성 요소만 미세 조정할 것입니다.
무엇을 동결(freeze)시킬지 선택하는 것은 무엇을 미세 조정할지 선택하는 것만큼 중요합니다.
LoRA는 low-rank adaptation의 약자로, 원래 가중치(original weights)를 동결한 상태로 작은 어댑터 행렬(adapter matrices)을 훈련하여 모델을 미세 조정합니다. 이는 학습 가능한 매개변수(trainable parameters)의 수와 그 기울기(gradients) 및 최적화기 상태(optimizer state)에 필요한 메모리를 줄여줍니다. 우리는 Hugging Face의 PEFT 라이브러리를 사용하여 이러한 어댑터를 추가하고 훈련합니다. 더 자세한 정보는 LoRA 가이드를 읽어보십시오.
Hugging Face Jobs는 클라우드 하드웨어에서 컨테이너화된 워크로드(containerized workloads)를 실행합니다. 우리는 이를 전처리 및 훈련에 사용합니다. 대신 자체 GPU를 사용하려면 아래 로컬 Docker 명령을 참조하십시오. 시작하기 전에 다음 사항들을 갖추었는지 확인하십시오:
- 여기서 얻을 수 있는 HF_TOKEN
- 귀하의 Hugging Face 계정은 DINOv3에 접근할 수 있어야 합니다. 왜냐하면 DINOv3는 게이티드 모델(gated model)이기 때문입니다.
- 자체 데이터셋을 사용하는 경우, 시작하기 전에 해당 메시(meshes)를 허브(Hub)에 업로드하십시오.
훈련하기 전에, TRELLIS.2가 어떻게 3D 객체를 생성하는지 그리고 우리가 LoRA 어댑터를 어디에 추가할 것인지 살펴보겠습니다.
TRELLIS.2는 다섯 가지 주요 단계(main stages)를 가지고 있습니다: 이미지 배경 제거(remove the image background), 이미지 인코딩(encode the image), 거친 구조 예측(predict the coarse structure), 상세 기하학 생성(generate detailed geometry), 그리고 재료 생성(generate materials). 이후 파이프라인은 결과를 텍스처가 입혀진 메시(textured mesh)로 내보냅니다.
먼저, BiRefNet이 입력 이미지에서 배경을 제거합니다. 이어서 DINOv3-L은 이미지를 특징 토큰(feature tokens)으로 인코딩합니다. 생성 모델들은 이 토큰들을 사용하여 3D 객체를 안내합니다.
희소 구조 흐름 DiT(sparse-structure flow DiT)는 객체의 거친 형태를 예측합니다. 이는 크기가 16³ × 8인 작은 잠재값(latent)을 생성하는 확산 변환기(diffusion transformer, DiT)입니다. 이후 변분 오토인코더(variational autoencoder, VAE) 디코더가 이 잠재값을 64 × 64 × 64의 점유 그리드(occupancy grid)로 변경합니다. 그리드의 각 복셀(voxel)은 해당 위치가 객체의 일부인지 여부를 보여줍니다. 복셀은 세 차원에서의 픽셀입니다. 우리는 이 모델에 LoRA 어댑터를 추가하는데, 이는 그 출력이 기하학 단계(geometry stage)를 안내하기 때문입니다.Stage 1
O-Voxel은 TRELLIS.2의 3D 형식입니다. 이는 희소 복셀 그리드(sparse voxel grid)입니다. 객체 표면에 있는 복셀만 유지합니다. 각 복셀은 메쉬 정점(mesh vertex)과 같은 국부 기하학 정보 및 재질 속성(material attributes)을 보존합니다. 이 형식이 작기 때문에, CPU는 몇 초 만에 메쉬를 O-Voxel로 변환할 수 있습니다. GPU는 몇 초 만에 O-Voxel을 다시 메쉬로 변환할 수 있습니다.
구조화된 희소 잠재값(Structured sparse latents)은 TRELLIS에서 나옵니다. 이는 3D 데이터를 희소 텐서(sparse tensors)에 보존하여 계산 속도를 높입니다. 이 모델은 그리드 내 복셀의 약 10%만 사용하므로, 빈 복셀에는 연산을 할애하지 않습니다. 따라서 형태 흐름 DiT(shape flow DiT)와 질감 흐름 DiT(texture flow DiT)는 희소 변환기(sparse transformers)입니다. 반면, 희소 구조 흐름 DiT는 밀집 변환기(dense transformer)입니다.O-Voxel
형태 흐름 DiT는 점유 그리드와 DINOv3 특징 토큰(feature tokens)을 사용하여 잠재값을 만듭니다. 또 다른 VAE 디코더가 이 잠재값을 해상도 512 × 512 × 512의 O-Voxel 기하학으로 변경합니다.
우리는 이 단계에서 DiT에 LoRA 어댑터를 추가하는데, 이는 객체의 형태를 제어하기 때문입니다.
마지막 단계에서 질감 흐름 DiT는 형태 잠재값과 DINOv3 특징을 입력으로 사용합니다. 이는 기하학 표면의 재질(materials)과 질감(textures)을 정의합니다. VAE 디코더가 이 잠재값을 각 복셀에 대한 기본 색상(base color), 금속성(metallic), 거칠기(roughness), 그리고 알파 값(alpha values)으로 변경합니다. 우리는 이 단계에는 LoRA 어댑터를 추가하지 않습니다. 우리의 데이터셋은 확산 색상 맵(diffuse color maps)을 가지고 있지만, 이 단계의 학습에 필요한 별도의 재질 맵은 부족합니다. 또한, 우리의 문제는 신발의 형태입니다. 기본 모델이 이미 좋은 색상과 질감을 구현하고 있습니다.
마지막으로, 파이프라인은 기하학적(geometry) 속성과 재질(material) 속성을 결합하여 텍스처가 입혀진 메시(textured mesh)를 만들고 이를 GLB 파일로 내보냅니다.
우리는 LoRA 어댑터를 희소 구조 흐름 DiT(sparse-structure flow DiT)와 형태 흐름 DiT(shape flow DiT)에만 추가합니다. 이 DiT들은 객체의 모양과 기하학을 제어합니다. 현재 기능을 유지하기 위해 텍스처 흐름 DiT(texture flow DiT)는 동결(freeze)시킵니다.
저희 데이터셋에는 베이크된 조명 효과(baked lighting effects)가 포함된 확산 색상 맵(diffuse color maps)이 있습니다. 텍스처 흐름 DiT를 학습시키려면 별도의 재질 맵(material maps)이 필요합니다. 저희 데이터셋의 색상 맵은 표면 색상과 조명을 혼합하고 있습니다. 따라서 이들은 이번 학습을 위한 정확한 타겟이 아닙니다.
만약 우리가 이 맵들로 텍스처 흐름 DiT를 학습시킨다면, 이는 조명을 표면 속성으로 표현하는 것을 배울 수 있습니다. 또한, 이는 치명적 망각(catastrophic forgetting)이라고 불리는 문제로 인해 이전에 배웠던 내용을 잊어버릴 수도 있습니다.
저희의 서브셋은 GSO에 있는 1,030개의 객체 중 188개의 신발을 포함하고 있습니다. 이렇게 작은 데이터셋의 경우, LoRA가 전체 파인튜닝(full fine-tuning)보다 더 나은 선택입니다. 전체 파인튜닝은 비용이 훨씬 많이 듭니다. 또한 모델이 이미 잘 수행하는 것(학습된 사전 지식, learned priors)을 덮어쓸 수 있습니다. 저희 구성은 두 개의 대상 DiT에 대해 총 약 930만개의 어댑터 매개변수(adapter parameters)를 학습시킵니다. 우리는 희소 구조 흐름 DiT의 어텐션 및 피드포워드(MLP) 레이어에 순위(rank) 8 어댑터를 사용합니다. 형태 흐름 DiT의 어텐션 레이어에는 순위 16 어댑터를 사용합니다. 이 순위는 어댑터 행렬의 크기를 설정합니다. 순위가 높을수록 어댑터가 더 많이 학습할 수 있게 합니다.
3D 모델 설정은 많은 종속성(dependencies)과 설치 단계가 필요할 수 있습니다. 저희는 필요한 소프트웨어가 포함된 두 개의 Docker 이미지를 제공합니다(Jobs 또는 자체 머신에서 실행 가능). 한 이미지는 전처리(preprocessing)용이고, 다른 하나는 파인튜닝(fine-tuning)용입니다.
LoRA 어댑터를 학습시키고 허브(Hub)에 업로드하려면:
- 컨테이너에 데이터 접근 권한을 부여합니다.
- 작업에 필요한 매개변수(parameters)를 설정합니다.
- 전처리 컨테이너를 실행합니다.
- 파인튜닝 컨테이너를 실행합니다.
그런 다음 저희는 파인튜닝 스크립트(fine-tuning script), 그 설정들, 그리고 각 설정이 학습에 어떻게 영향을 미치는지 설명합니다.
실험의 세 부분에 대해 안내하겠습니다:
데이터 준비(Prepare the data). GLB 또는 OBJ 메시를 TRELLIS.2가 학습에 사용하는 잠재 표현(latent representations)으로 변환하고 입력 이미지를 렌더링합니다.
어댑터 학습(Train the adapters). 준비된 데이터셋과 LoRA 구성을 사용하여 각 DiT에 대해 하나의 작업을 실행합니다.
코드 설명(Explain the code). 컨테이너가 데이터를 로드하고, 어댑터를 추가하며, 모델을 학습하는 방식을 검토합니다.
모든 학습 실행이 재사용할 수 있도록 전처리된 데이터셋을 저장합니다.
전처리 컨테이너는 Blender를 사용하여 각 객체의 16개 뷰(views)를 렌더링합니다. LoRA 학습에서는 이 뷰들을 조건화 이미지(conditioning images)로 사용하게 됩니다. Blender는 xvfb-run 유틸리티를 통해 디스플레이 없이 실행됩니다. 또한, 컨테이너는 모든 객체를 전처리 파이프라인을 거쳐 처리합니다. 중간 출력 예시로는 여기의 신발 데이터셋을 참고하십시오. 저희가 다룬 188개 객체의 전체 실행은 a100-large 작업에서 약 1시간이 소요되었습니다. 전처리 파이프라인의 간소화된 버전도 이 보조 저장소(companion repository)에서 이용할 수 있습니다.
export HF_TOKEN=hf_xxx
# Replace your-username with your Hugging Face username
hf jobs run -d --flavor a100-large --timeout 6h --secrets HF_TOKEN \
...
Hugging Face Jobs의 경우, 다음과 같은 인수를 전달합니다:
| 플래그 | 의미 |
|---|
-d | Detach: 작업 ID와 URL을 출력한 후 반환됩니다. 작업을 추적하려면 hf jobs logs -f <id>를 실행하십시오. |
--flavor a100-large | 하드웨어 사양입니다. 저희는 a100-large (80 GB)를 사용합니다. 모든 옵션을 보려면 hf jobs hardware를 실행하십시오. |
--timeout 6h | 이 시간이 지나면 작업이 중지됩니다. 두 명령어 모두 재개할 수 있습니다. 전처리는 완료된 객체를 건너뜁니다. 학습은 --push-to에서 가장 최신 체크포인트부터 계속됩니다. |
--secrets HF_TOKEN | 토큰(HF_TOKEN 또는 hf auth login에서 가져옴)을 암호화된 환경 변수로 작업에 전송합니다. 이 토큰은 이미지에 기록되지 않습니다. |
cr.hf.co/...:latest | 실행할 이미지입니다. 그 뒤의 모든 텍스트는 컨테이너 내에서 실행되는 명령어입니다. |
이 컨테이너는 (Hugging Face Jobs에서 실행하든 Docker로 실행하든 관계없이) 몇 가지 추가 인수를 받으며, 자세한 내용은 여기에서 확인할 수 있습니다. 전처리 파이프라인에서 수행되는 단계들은 본 섹션의 나중에 문서화되어 있습니다.
학습 컨테이너는 준비된 데이터셋과 구성 프리셋 또는 파일을 로드합니다. 이 컨테이너는 Trackio를 사용하여 학습 메트릭을 기록합니다. 실행 모니터링에는 Trackio 대시보드를 사용하십시오.
저희는 LoRA(Low-Rank Adaptation) 학습을 위한 두 가지 모드를 나타내기 위해 ‘ss’와 ‘shape’를 사용합니다. ‘ss’ 모드는 Sparse-Structure Flow DiT만 파인튜닝하는 반면, ‘shape’ 모드는 Shape Flow DiT를 파인튜닝합니다. 내부적으로 이들은 각 단계를 위해 사용할 다른 config.json 파일을 참조할 뿐입니다. 프로세스에 대한 더 자세한 내용은 본 섹션의 나중에 확인할 수 있습니다.
export HF_TOKEN=hf_xxx # Replace your-username with your Hugging Face username
hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
...
만약 --push-to를 설정하면, 컨테이너는 어댑터 체크포인트를 Hub에 업로드합니다. 이 컨테이너는 간결함을 위해 여기에서도 문서화된 매개변수를 받습니다. 또한 동일한 저장소에서 간편 스크립트도 사용할 수 있습니다.
이미 처리된 데이터셋으로 테스트하거나 파인튜닝을 하고 싶다면, 다음 명령어 세트를 자유롭게 사용하십시오.
export HF_TOKEN=hf_xxx
# Replace your-username with your Hugging Face username
# For the Sparse-Structure Flow DiT
...
명령어와 플래그는 Hugging Face Jobs의 것과 동일합니다. 다를 뿐인 것은 Docker 옵션입니다.
# For preprocessing into latents
docker run --gpus all --rm --shm-size 16g -e HF_TOKEN=hf_xxx \
-v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
...
# For Sparse-structure Flow DiT
docker run --gpus all --rm --shm-size 16g \
-v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
...
이미 사전 처리된 데이터셋을 사용자의 선택한 로컬 머신에서 사용하려면 다음 명령어를 사용하세요:
export HF_TOKEN=hf_xxx # your-username을 Hugging Face 사용자 이름으로 교체하십시오
# Sparse-structure Flow DiT의 경우
docker run --gpus all --rm --shm-size 16g -e HF_TOKEN \
...
저희가 학습시킨 LoRA 어댑터는 이 저장소에 있습니다: suvadityamuk/trellis-2-finetune-gso
Preprocessing
저희는 또한 Blender에서 각 객체의 16개 뷰를 렌더링했습니다. 학습 중에는 각 샘플에 대해 무작위로 선택된 하나의 뷰를 컨디셔닝 이미지(conditioning image)로 사용합니다. 저희는 188개의 객체 중 9개를 평가용으로 보관하고 나머지 179개로 학습을 진행했습니다. 이 제외된(held-out) 객체들은 어댑터가 모델이 학습 과정에서 본 적 없는 객체에서도 작동하는지 보여줍니다.
Blender가 하나의 객체에 대해 렌더링한 일부 뷰들
첫 번째 스니펫은 데이터를 설정합니다. FLOW는 학습할 흐름(flow) DiT를 선택하며 (ss 또는 shape), 그리고 STEPS가 있습니다.
