คู่มือ Fine-Tuning โมเดล TRELLIS.2 สำหรับงานสร้างภาพ 3D

TLDR; บล็อกนี้จะนำเสนอสคริปต์และคอนเทนเนอร์สำหรับการ Fine-tuning โมเดล TRELLIS.2 ซึ่งเป็นโมเดลสร้างภาพ 3D ยอดนิยมบนข้อมูลส่วนตัวของคุณ โดยสามารถรันได้ทั้งแบบเครื่องโลคอลและบน Hugging Face Jobs หากคุณสนใจสามารถ ใช้สคริปต์ของเรา โดยระบุ argument เพียงไม่กี่ตัว หรือจะลองปรับแต่งด้วยตัวเอง ผ่าน repository นี้ ส่วนใครที่อยากเจาะลึก วิธีการทำงานของสคริปต์ เรามีรายละเอียดให้ติดตามกันครับ

Introduction

การสร้างโมเดล 3D มีความสำคัญอย่างมากในงานหุ่นยนต์ การจำลอง (simulation) และสื่อบันเทิง แต่ความท้าทายคือการทำให้วัตถุที่สร้างขึ้นตรงตามกฎเกณฑ์ที่เข้มงวดและมีมาตรฐานเดียวกัน เนื่องจากปัจจุบันมีฟอร์แมตไฟล์ที่หลากหลาย เช่น GLB และ PLY รวมถึง การแสดงผลข้อมูล (representations) หลายรูปแบบ อาทิ neural radiance fields (NeRFs) และ 3D Gaussian splatting (3DGS)

TRELLIS เข้ามาแก้ปัญหานี้ด้วยการนำเสนอ unified latent representation พร้อม decoder ที่สามารถแปลงการแสดงผลนี้เป็นไฟล์ 3D ในฟอร์แมตที่ต้องการได้ทันที ต่อมาทีมงานเดิมได้เปิดตัว TRELLIS.2 ที่ปรับปรุงโครงสร้างให้ดีขึ้น อย่างไรก็ตาม ทั้งสองโมเดลยังขาดวิธีที่ง่ายในการปรับแต่ง (adapt) ให้เข้ากับงานเฉพาะเจาะจง

สิ่งที่เราต้องการปรับปรุง

โมเดลอเนกประสงค์อย่าง TRELLIS.2 อาจผิดพลาดในรายละเอียดสำคัญของงานเฉพาะทาง เช่น เรขาคณิต (geometry) ของเฟอร์นิเจอร์ที่ต้องการความแม่นยำสูง เราจึงทดลองใช้ชุดข้อมูล GSO หมวดรองเท้า 188 ชิ้นเพื่อทดสอบว่า LoRA จะช่วยให้โมเดลสร้างพื้นรองเท้าได้ถูกต้องขึ้นหรือไม่ จากเดิมที่โมเดล pretrained มักสร้างพื้นรองเท้าผิดเพี้ยน

เป้าหมายคือการใช้ชุดข้อมูล Google Scanned Objects (GSO) เพื่อเพิ่มประสิทธิภาพให้ TRELLIS.2 โดยเราจะอธิบายว่าส่วนประกอบใดใน pipeline ที่ควร Fine-tune และส่วนใดควรแช่แข็ง (freeze) ไว้ เพื่อให้การปรับแต่งเน้นไปที่จุดสำคัญของงานเท่านั้น

การเลือกสิ่งที่จะแช่แข็ง (freeze) ไว้มีความสำคัญพอๆ กับการเลือกสิ่งที่จะ Fine-tune

เราจะใช้อะไรบ้าง?

LoRA

LoRA หรือ low-rank adaptation คือการฝึกเมทริกซ์ adapter ขนาดเล็กในขณะที่แช่แข็งน้ำหนัก (weights) เดิมของโมเดลไว้ ช่วยลดการใช้หน่วยความจำและจำนวนพารามิเตอร์ที่ต้องฝึก โดยเราจะใช้ไลบรารี PEFT ของ Hugging Face เข้ามาช่วยในส่วนนี้ ศึกษาเพิ่มเติมได้ที่ คู่มือ LoRA

Hugging Face Jobs

Hugging Face Jobs ช่วยรันเวิร์กโหลดแบบ containerized บนคลาวด์ เหมาะสำหรับการทำ preprocessing และการฝึกโมเดล สิ่งที่คุณต้องเตรียมคือ:

• HF_TOKEN รับได้ จากที่นี่
• สิทธิ์เข้าถึง DINOv3 ซึ่งเป็น gated model
• อัปโหลดไฟล์ meshes ของคุณไปยัง Hub ก่อนเริ่ม

TRELLIS.2 ทำงานอย่างไร

TRELLIS.2 มีขั้นตอนหลัก 5 ขั้น: ลบพื้นหลังภาพ, เข้ารหัสภาพ, ทำนายโครงสร้างคร่าวๆ, สร้างเรขาคณิตรายละเอียดสูง และสร้างวัสดุ ก่อนส่งออกเป็น textured mesh

การเข้ารหัสภาพอินพุต

BiRefNet จะลบพื้นหลังอินพุต แล้วใช้ DINOv3-L เข้ารหัสภาพเป็น feature tokens เพื่อนำทาง (guide) การสร้างวัตถุ 3D

การสร้างโครงสร้างและเรขาคณิต

โมเดลใช้ sparse-structure flow DiT ทำนายรูปทรงคร่าวๆ เป็น occupancy grid ขนาด 64³ ซึ่งเราจะเพิ่ม LoRA adapter ในจุดนี้ จากนั้นใช้ shape flow DiT เพื่อสร้างเรขาคณิตโดยละเอียดในรูปแบบ O-Voxel (ความละเอียด 512³) ซึ่งจุดนี้เป็นอีกจุดที่เราจะเพิ่ม LoRA เพราะมีผลต่อรูปทรงวัตถุโดยตรง

การสร้างวัสดุและ Mesh

ในขั้นตอนสุดท้าย texture flow DiT จะกำหนดวัสดุและพื้นผิว แต่ในการทดลองนี้เราจะ ไม่เพิ่ม LoRA adapter ในส่วนนี้ เนื่องจากปัญหาหลักอยู่ที่รูปทรงรองเท้า และโมเดลพื้นฐานทำเรื่องสีและพื้นผิวได้ดีอยู่แล้ว อีกทั้งชุดข้อมูลของเราไม่มีแผนที่วัสดุแยกต่างหากที่จำเป็นต่อการฝึก

Planning และการรันระบบ

เราเลือก Fine-tune เฉพาะส่วนที่ควบคุมรูปทรงเพื่อป้องกันปัญหา catastrophic forgetting และใช้ LoRA เพราะชุดข้อมูลมีขนาดเล็ก (รองเท้า 188 คู่) ซึ่งประหยัดค่าใช้จ่ายและรักษาความสามารถเดิมของโมเดลได้ดีกว่า

วิธีรันบน Hugging Face Jobs

ขั้นตอน Preprocessing:

export HF_TOKEN=hf_xxx  
hf jobs run -d --flavor a100-large --timeout 6h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-prep:latest \
  trellis2-prep \
    --meshes your-username/my-meshes \
    --push-to your-username/my-meshes-trellis2 \
    --private --holdout 0.05 \
    --views 16 --workers 8

ขั้นตอน LoRA Fine-tuning: รันแยกกันสำหรับส่วนโครงสร้าง (ss) และเรขาคณิต (shape)

hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train --config ss \
    --latents your-username/my-meshes-trellis2 \
    --push-to your-username/my-trellis2-lora \
    --run-name ss_lora \
    --trackio-space your-username/trellis2-trackio

ผลลัพธ์ (Results)

การ Fine-tune ช่วยให้โมเดลสร้างพื้นรองเท้าได้ใกล้เคียงกับของจริงมากขึ้นอย่างเห็นได้ชัด โดยวัดจากเกณฑ์ F-score ที่เพิ่มขึ้นจาก 0.276 เป็น 0.803 และค่า Chamfer distance ที่ลดลงอย่างมาก

บทเรียนสำคัญคือการเลือกส่วนที่จะ Fine-tune ให้สอดคล้องกับปัญหา ในอนาคตเราสามารถนำวิธีการนี้ไปประยุกต์ใช้กับโมเดลอื่นๆ เช่น Pixal3D หรือโมเดลสร้างโลกอย่าง HunyuanWorld 2.0 ได้เช่นกัน