Tái hiện Stable Diffusion 3.5 từ đầu bằng PyTorch thuần túy

(github.com/yousef-rafat)

2 điểm bởi GN⁺ 2025-06-15 | 1 bình luận | Chia sẻ qua WhatsApp

Dự án miniDiffusion là một mã nguồn mở tái hiện mô hình Stable Diffusion 3.5 từ đầu chỉ bằng PyTorch
Cấu trúc của dự án này nổi bật ở chỗ tập trung vào mục đích giáo dục và phục vụ thử nghiệm, hack
Toàn bộ codebase chỉ khoảng 2800 dòng, được cấu thành bằng lượng mã tối thiểu, từ VAE đến DiT, cùng các script huấn luyện và dataset
Các thành phần chính gồm VAE, bộ mã hóa văn bản CLIP, T5, transformer khuếch tán đa phương thức, joint attention
Hiện vẫn bao gồm các tính năng thử nghiệm và cần thêm nhiều kiểm thử

Giới thiệu dự án miniDiffusion

miniDiffusion là một dự án mã nguồn mở tái hiện các chức năng cốt lõi của Stable Diffusion 3.5 chỉ bằng PyTorch
So với Stable Diffusion 3.5 hiện có, dự án này có những ưu điểm sau

Codebase chỉ khoảng 2.800 dòng, rất nhỏ gọn nên rất phù hợp để trực tiếp phân tích cấu trúc và học tập
Có thể được tận dụng hữu ích cho nhiều thử nghiệm machine learning và hack mô hình
Có rất ít phụ thuộc và chỉ sử dụng số lượng thư viện tối thiểu

Cấu trúc cốt lõi và các tệp thành phần

dit.py : phần hiện thực mô hình Stable Diffusion chính
dit_components.py : các thành phần như embedding, chuẩn hóa, patch embedding và hàm phụ trợ cho DiT
attention.py : phần hiện thực thuật toán Joint Attention
noise.py : bao gồm bộ lập lịch Euler ODE cho Rectified Flow
t5_encoder.py, clip.py : hiện thực bộ mã hóa văn bản T5 và CLIP
tokenizer.py : hiện thực tokenizer Byte-Pair và Unigram
metrics.py : hiện thực chỉ số đánh giá FID (Fréchet inception distance)
common.py : cung cấp các hàm phụ trợ cần cho huấn luyện
common_ds.py : hiện thực dataset iterable chuyển đổi hình ảnh thành dữ liệu huấn luyện cho DiT
thư mục model : lưu checkpoint mô hình và log sau huấn luyện
thư mục encoders : lưu checkpoint của các mô-đun riêng như VAE, CLIP

⚠️ Tính năng thử nghiệm và nhu cầu kiểm thử miniDiffusion hiện vẫn bao gồm các tính năng thử nghiệm và cần thêm nhiều kiểm thử

Cấu hình chi tiết theo từng tính năng chính

Core Image Generation Modules

Hiện thực VAE, bộ mã hóa văn bản CLIP, T5
Hiện thực tokenizer Byte-Pair, Unigram

SD3 Components

Multi-Modal Diffusion Transformer Model
Hiện thực Flow-Matching Euler Scheduler
Logit-Normal Sampling
Áp dụng thuật toán Joint Attention

Script huấn luyện và suy luận mô hình

Cung cấp script huấn luyện và suy luận cho SD3 (Stable Diffusion 3.5)

Giấy phép

Được phát hành theo giấy phép MIT và được tạo ra cho mục đích giáo dục và thử nghiệm

Ý nghĩa và ưu điểm của dự án mã nguồn mở này

Có thể trực tiếp huấn luyện và hack cấu trúc mô hình tạo ảnh hiện đại ở cấp độ Stable Diffusion 3.5 chỉ bằng PyTorch thuần túy
Mã nguồn ngắn gọn và độc lập, được tối ưu cho phân tích cấu trúc / tinh chỉnh mô hình / nghiên cứu thuật toán mới
Có thể trực tiếp thực hành các kỹ thuật đa phương thức, transformer, attention hiện đại
Cung cấp nền tảng để thử nghiệm an toàn, tách biệt với các dự án thương mại

1 bình luận

GN⁺ 2025-06-15

Ý kiến trên Hacker News

Bản triển khai tham chiếu của Flux thực sự có cấu trúc rất tối giản, nên nếu ai quan tâm thì rất đáng xem thử
- GitHub Flux
- Dự án minRF có ưu điểm là dùng rectified flow nên khá dễ để bắt đầu khi huấn luyện các mô hình diffusion nhỏ
- GitHub minRF
- Bản triển khai tham chiếu của Stable Diffusion 3.5 cũng được viết khá gọn gàng, phù hợp để tham khảo
- GitHub SD 3.5
- Các bản triển khai tham chiếu thường không được bảo trì tốt và có khá nhiều lỗi
  - Ví dụ có thể kể đến tokenizer CLIP của OpenAI; đây không phải bản đã dùng trong huấn luyện chính thức mà chỉ là bản tham chiếu, nhưng lỗi vẫn không được sửa và lại bị sao chép nguyên sang nhiều dự án khác
  - Flux cũng không phải thứ thực sự đã được dùng để huấn luyện, và có lỗi có thể gây ra một số vấn đề với cudagraphs, v.v.
  - Nó phụ thuộc vào tham chiếu CLIP, mà bản thân CLIP lại có lỗi, nên lỗi tiếp tục bị lan truyền theo kiểu dây chuyền
Tôi tự hỏi liệu có phải dự án miniDiffusion đang dùng mô hình Stable Diffusion 3.5 hay không
- Đoạn mã liên quan
- Bộ dữ liệu huấn luyện rất nhỏ và chỉ gồm ảnh liên quan đến thời trang
- Bộ dữ liệu thời trang
- Bộ dữ liệu đó dùng để thực hành fine-tune mô hình diffusion
  - Đây là bản tái triển khai SD3 bằng mã mới, nhưng do giới hạn phần cứng của tôi nên trọng số được lấy từ HuggingFace để dùng
Tôi tò mò liệu dùng PyTorch thuần có mang lại lợi thế hiệu năng trên GPU không phải của NVIDIA hay không, hay PyTorch đã được tối ưu cho CUDA quá mạnh đến mức các hãng GPU khác không thể cạnh tranh
- PyTorch chạy khá ổn trên Apple Silicon
  - Tuy nhiên, GPU của Apple khó đạt hiệu năng tính toán ngang với các GPU NVIDIA đầu bảng, nên khó so sánh trực tiếp
  - Nhân tiện, khi dùng PyTorch trên Apple Silicon có một điểm hơi đặc biệt
    - Mỗi tensor được xem là thuộc về một thiết bị cụ thể (CPU hoặc GPU), nên khi di chuyển dữ liệu sẽ phát sinh sao chép toàn bộ
    - Mac dùng kiến trúc bộ nhớ hợp nhất, nhưng PyTorch vẫn hoạt động theo cách phải sao chép dữ liệu
- Trên các thiết bị không phải NVIDIA như AMD, vẫn có thể chạy workload ML qua Vulkan
  - Nếu các phần mở rộng mới như cooperative matrix và hỗ trợ tính năng mới ở cấp driver tiếp tục được bổ sung, thì khoảng cách hiệu năng chỉ vài phần trăm so với CUDA cũng có thể biến mất
- Hỗ trợ ROCm của PyTorch tiến triển rất chậm, và kể cả khi chạy được thì tốc độ cũng chậm
- PyTorch đúng là có chạy được trên ROCm, nhưng tôi không rõ nó có thực sự tốt đến mức hoàn toàn "ngang hàng" hay không

Trong mã PyTorch, thay vì

    self.q = nn.Linear(embed_size, embed_size, bias = False)
    self.k = nn.Linear(embed_size, embed_size, bias = False)
    self.v = nn.Linear(embed_size, embed_size, bias = False)

có đề xuất rằng nên thử

    self.qkv = nn.Linear(embed_size, 3*embed_size, bias = False)
    # bên trong hàm forward
    qkv = self.qkv(x)

như vậy

Làm như vậy thì thay vì các tham số của q, k, v ban đầu được nối độc lập, các tham số giữa q, k, v sẽ được nối với nhau
- Nếu tôi đang quá mệt nên nhầm lẫn thì mong mọi người thông cảm
Có vẻ đây là tài liệu tốt cho người học
- Tôi tự hỏi liệu có tutorial hoặc tài liệu giải thích nào mà người mới bắt đầu cũng có thể làm theo không
- Có khóa học của fast.ai hướng dẫn tự triển khai Stable Diffusion
  - Khóa học fast.ai
Tôi thắc mắc liệu điều này có nghĩa là có thể dùng Stable Diffusion mà không bị ràng buộc bởi giấy phép hay không
- Không phải vậy; thuật toán suy luận/huấn luyện (bản thân phần toán học) không thuộc đối tượng bản quyền, nhưng OP chỉ đơn giản là tái triển khai phần mã
  - Thứ liên quan đến bản quyền là bản thân mô hình (trọng số), và OP không tự huấn luyện vì không có dữ liệu hay sức mạnh tính toán
Thành thật mà nói thì hơi ngại, nhưng tôi muốn biết chúng ta thực sự có thêm được điều gì mới trước và sau khi kho mã này xuất hiện
- Cá nhân tôi vẫn tránh tự làm mô hình và chủ yếu chỉ đứng ngoài quan sát kết quả
- Tôi vốn mơ hồ cho rằng từ trước đã có sẵn script suy luận/huấn luyện dựa trên PyTorch được công khai rồi
- Ít nhất tôi nghĩ script suy luận sẽ đi kèm khi phát hành mô hình, và chắc cũng sẽ có script fine-tune/huấn luyện
- Tôi không chắc dự án này là kiểu "clean room" hay "dirty room" viết lại từ cái có sẵn, hay là vì ngay cả mã PyTorch hiện có cũng quá rối do dính CUDA/C nên một bản PyTorch thuần mới thực sự có ý nghĩa lớn
- Dù sao tôi cũng không rõ, nên sẽ rất tốt nếu ai đó có thể giải thích
  - Giá trị cốt lõi của dự án này là một bản triển khai "ít phụ thuộc nhất có thể"
    - Tôi chưa từng trực tiếp chạy SD 3.5, nhưng nó được xây trên thư viện huggingface, mà cá nhân tôi thấy huggingface có hệ phụ thuộc quá phức tạp, đến mức nếu môi trường không gần như y hệt môi trường của tác giả thì rất khó chạy nổi
    - Đặc biệt là sau vài tháng hay vài năm kể từ bản phát hành gốc, việc chạy được một mô hình cụ thể thường trở nên rất khó
    - Ví dụ, nếu xem file requirements.txt của bản tham chiếu Stability AI cho SD3.5, thậm chí còn không ghim phiên bản và lại bao gồm các thư viện khổng lồ như transformers, nên trong môi trường làm việc thực tế điều này thực sự rất khó chịu
  - Stability AI phát hành các mô hình Stable Diffusion theo Stability AI Community License, nên không phải là "hoàn toàn tự do" như MIT
    - Không được phép sửa đổi trọng số theo một số cách nhất định
    - Gói này được thiết kế để có thể chạy mô hình (suy luận), hoặc fine-tune bằng các trọng số AI đã có sẵn
    - Nó rất tuyệt cho mục đích học tập/nghiên cứu, nhưng vấn đề giấy phép vẫn còn tồn tại
Khi nghĩ về SD 3.5, hoặc bất kỳ phiên bản nào, tôi xem phần cốt lõi nằm ở các trọng số được tạo ra trong quá trình huấn luyện
- Bản thân mã có tầm quan trọng tương đối thấp hơn về mặt chất lượng kết quả hay hiệu năng
- Nhưng đây có thể không phải nhận định chính xác, và tôi không có ý hạ thấp những nỗ lực như thế này
Tôi tò mò về mức độ có thể dùng trong thực tế của mã nguồn học thuật gốc do nhóm CompViz của Ludwig Maximilian University công bố
Tôi muốn biết liệu phần triển khai diffusion transformer (DiT) ở đây có thực sự triển khai đúng cross-token attention như bản SD 3.5 đầy đủ hay không, hay đã được đơn giản hóa để mã dễ đọc hơn