Stable Diffusion 3.5을 순수 PyTorch로 처음부터 다시 구현함

18 hours ago 1

  • miniDiffusion 프로젝트는 Stable Diffusion 3.5 모델을 PyTorch만을 사용하여 처음부터 재구현한 오픈소스임
  • 이 프로젝트의 구조는 교육 목적실험, 해킹 용도에 초점을 둔 것이 특징임
  • 전체 코드베이스는 약 2800줄로, VAE부터 DiT, 학습 및 데이터셋 스크립트까지 최소한의 코드로 구성됨
  • 주요 구성 요소로는 VAE, CLIP, T5 텍스트 인코더, 멀티모달 디퓨전 트랜스포머, 공동 어텐션 등이 있음
  • 아직 실험적인 기능이 포함된 상태로서, 더 많은 테스트가 필요한 상태임

miniDiffusion 프로젝트 소개

miniDiffusionStable Diffusion 3.5의 핵심 기능을 PyTorch만으로 재구현한 오픈소스 프로젝트임
이 프로젝트는 기존 Stable Diffusion 3.5와 비교해 다음과 같은 장점이 있음

  • 코드베이스가 약 2,800줄로 크기가 작아, 직접 구조를 분석하고 학습하기에 매우 적합
  • 다양한 기계학습 실험모델 해킹에 유용하게 활용 가능함
  • 종속성이 매우 적으며 최소한의 라이브러리만을 사용함

핵심 구조 및 구성 파일

  • dit.py : 메인 Stable Diffusion 모델 구현부
  • dit_components.py : 임베딩, 정규화, 패치 임베딩, DiT 보조 함수 구성
  • attention.py : Joint Attention(공동 어텐션) 알고리듬 구현부
  • noise.py : Rectified Flow 를 위한 Euler ODE 스케줄러 포함
  • t5_encoder.py, clip.py : T5 및 CLIP 텍스트 인코더 구현
  • tokenizer.py : Byte-Pair 및 Unigram 토크나이저 구현
  • metrics.py : FID(Fréchet inception distance) 평가 지표 구현
  • common.py : 학습에 필요한 보조 함수 제공
  • common_ds.py : 이미지를 DiT용 학습 데이터로 변환하는 iterable 데이터셋 구현
  • model 폴더 : 학습 이후 모델 체크포인트와 로그 저장
  • encoders 폴더 : VAE, CLIP 등 별도 모듈의 체크포인트 저장

⚠️ 실험적 기능 및 테스트 필요성 miniDiffusion은 아직 실험적인 기능들을 포함하고 있으며, 더 많은 테스트가 필요한 상태임

주요 기능별 세부 구성

Core Image Generation Modules

  • VAE, CLIP, T5 텍스트 인코더 구현
  • Byte-Pair, Unigram 토크나이저 구현

SD3 Components

  • Multi-Modal Diffusion Transformer Model
  • Flow-Matching Euler Scheduler 구현
  • Logit-Normal Sampling
  • Joint Attention 알고리듬 도입

모델 학습 및 추론 스크립트

  • SD3(Stable Diffusion 3.5)용 학습 및 추론 스크립트 제공

라이선스

  • MIT 라이선스로 공개되어 있으며, 교육 및 실험 목적으로 제작됨

이 오픈소스 프로젝트의 의미와 장점

  • Stable Diffusion 3.5 수준의 최신 이미지 생성 모델 구조를 순수 PyTorch만으로 직접 학습·해킹 가능함
  • 코드가 간결하고 독립적이라 구조 분석/모델 튜닝/신규 알고리듬 연구에 최적화되어 있음
  • 최신 멀티모달, 트랜스포머, 어텐션 기법 등을 직접 실습할 수 있음
  • 상업 프로젝트와 별개로 안전하게 실험할 수 있는 기반 제공

Read Entire Article