Trellis.2 trên Mac Silicon: Chạy mô hình chuyển đổi ảnh sang 3D mà không cần GPU Nvidia

Trellis.2 trên Mac Silicon: Chạy mô hình chuyển đổi ảnh sang 3D mà không cần GPU Nvidia

Một nhà phát triển có tên Shivampkumar đã thực hiện một bước tiến thú vị khi chuyển port mô hình TRELLIS.2 của Microsoft để chạy trực tiếp trên các máy Mac sử dụng chip Apple Silicon. Đây là một tin vui cho cộng đồng người dùng Mac, đặc biệt là những người làm việc trong lĩnh vực đồ họa 3D và AI, vì họ không còn cần sở hữu GPU đắt tiền của Nvidia để trải nghiệm công nghệ này.

Tổng quan về dự án

TRELLIS.2 là một mô hình chuyển đổi ảnh sang 3D (image-to-3D) tiên tiến với 4 tỷ tham số. Bản gốc của mô hình này yêu cầu CUDA và các thư viện chuyên dụng như flash_attn, nvdiffrast cùng các nhân tích chập thưa (sparse convolution kernels) tùy chỉnh — những thứ vốn không hoạt động trên macOS.

Để giải quyết vấn đề này, bản port mới sử dụng PyTorch MPS (Metal Performance Shaders) và thay thế hoàn toàn các thao tác riêng biệt của CUDA bằng các giải pháp thay thế thuần túy bằng PyTorch và Python.

Những thay đổi kỹ thuật chính

Thay vì viết lại toàn bộ, tác giả đã tập trung vào việc thay thế các thành phần phụ thuộc vào CUDA:

• Tích chập thưa 3D (Sparse 3D Convolution): Thay thế flex_gemm bằng phương pháp gather-scatter được viết bằng Python thuần túy.
• Attention: Sử dụng torch.nn.functional.scaled_dot_product_attention (SDPA) tích hợp sẵn của PyTorch thay vì flash_attn.
• Trích xuất lưới (Mesh Extraction): Viết lại thuật toán trích xuất lưới bằng Python dictionary thay vì các thao tác hashmap trên CUDA.

Tổng cộng, chỉ khoảng vài trăm dòng code trên 9 tệp đã được thay đổi để đạt được mục đích này.

Hiệu suất trên phần cứng Apple

Mặc dù không nhanh bằng các hệ thống sử dụng GPU H100 chuyên dụng, nhưng hiệu suất trên Mac Silicon vẫn rất ấn tượng đối với một thiết bị làm việc cá nhân:

• Thời gian tạo: Khoảng 3,5 phút để tạo một mô hình 3D từ một bức ảnh đơn trên chip M4 Pro (24GB RAM).
• Đầu ra: Tạo ra các tệp lưới (mesh) với khoảng 400.000 đỉnh (vertices).
• Bộ nhớ: Sử dụng đỉnh khoảng 18GB bộ nhớ thống nhất (unified memory) trong quá trình tạo.

Hạn chế hiện tại

Dù là một thành tựu lớn, bản port này vẫn có một số hạn chế so với bản gốc CUDA:

• Không xuất kết cấu (Texture export): Việc tạo kết cấu yêu cầu nvdiffrast (chỉ hỗ trợ CUDA). Do đó, các tệp xuất ra chỉ có màu đỉnh (vertex colors).
• Tốc độ: Tích chập thưa bằng PyTorch chậm hơn khoảng 10 lần so với nhân CUDA flex_gemm.
• Chỉ hỗ trợ suy luận (Inference only): Hiện tại dự án chỉ hỗ trợ tạo mô hình, không hỗ trợ đào tạo (training).

Cách cài đặt

Để chạy dự án này, bạn cần một máy Mac chạy chip Apple Silicon (M1 trở lên), Python 3.11+ và ít nhất 24GB RAM được khuyến nghị để xử lý mô hình 4B lớn.

Quy trình cài đặt khá đơn giản với việc clone repository từ GitHub, đăng nhập vào HuggingFace để lấy quyền truy cập các mô hình có khóa (gated models), và chạy script cài đặt tự động.

Dự án này mở ra cơ hội lớn cho các nhà phát triển và nhà sáng tạo nội dung sử dụng hệ sinh thái Apple, cho phép họ thử nghiệm các công nghệ AI 3D tiên tiến mà không cần phụ thuộc vào các máy chủ đám mây hay phần cứng GPU rời.

---

Nguồn: Shivampkumar (GitHub)