▲

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

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

miniDiffusion 프로젝트 소개

miniDiffusion는 Stable 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만으로 직접 학습·해킹 가능함
  • 코드가 간결하고 독립적이라 구조 분석/모델 튜닝/신규 알고리듬 연구에 최적화되어 있음
  • 최신 멀티모달, 트랜스포머, 어텐션 기법 등을 직접 실습할 수 있음
  • 상업 프로젝트와 별개로 안전하게 실험할 수 있는 기반 제공
▲
GN⁺ 1일전  [-]
Hacker News 의견

  • Flux 레퍼런스 구현은 정말 미니멀한 구조이기 때문에 관심 있는 사람이 있다면 한번 살펴볼 만한 가치가 있음

    • Flux 깃허브

    • minRF 프로젝트는 정정된 플로우(rectified flow)를 활용해서 작은 디퓨전 모델을 학습할 때 쉽게 시작할 수 있는 점이 장점임

    • minRF 깃허브

    • Stable Diffusion 3.5의 레퍼런스 구현도 상당히 간결하게 짜여있어서 참고하기에 적합함

    • SD 3.5 깃허브

    • 레퍼런스 구현들은 관리가 잘 안 되고 버그가 많은 경우가 많음

      • 예시로 OpenAI의 CLIP 토크나이저 같은 경우를 들 수 있는데, 공식 학습에 쓰인 버전은 아니고 레퍼런스일 뿐인데도 버그가 수정되지 않은 채 여러 프로젝트에 똑같이 복제되는 상황임
      • Flux도 학습에 실제로 사용된 건 아니고, cudagraphs 등에서 약간 문제를 일으킬 수 있는 버그가 있음
      • CLIP 레퍼런스에 의존하는데 CLIP 자체가 버그가 있으니, 마찬가지로 버그가 전파되는 구조임

  • miniDiffusion 프로젝트가 Stable Diffusion 3.5 모델을 사용한다는 의미인지 궁금증이 생김

    • 관련 코드

    • 학습 데이터셋이 매우 작고 패션 관련 사진만 포함

    • 패션 데이터셋

    • 해당 데이터셋은 디퓨전 모델의 파인튜닝을 실습해보기 위한 용도임

      • SD3를 새로 코드로 재구현한 것이지만, 내 하드웨어의 한계로 인해 가중치는 HuggingFace에서 가져와서 사용함

  • 순수 PyTorch를 사용하면 NVIDIA가 아닌 GPU에서 성능 이점이 생기는지, 혹은 PyTorch가 CUDA에 워낙 최적화되어 있어서 다른 GPU 벤더는 경쟁이 불가능한지 궁금함

    • PyTorch는 Apple Silicon에서도 꽤 잘 작동하는 편임

      • 다만 애플 GPU는 NVIDIA 최상급 GPU 수준의 컴퓨팅 성능을 내기는 어렵기 때문에 직접 비교는 어려움
      • 참고로, 애플 실리콘에서 PyTorch를 사용할 때 약간 특이점이 있음
        • 각 텐서가 특정 디바이스(CPU나 GPU)에 '소유'된 것으로 인식되기 때문에, 데이터 이동 시 전체 복사가 발생
        • 맥은 통합 메모리 구조지만 PyTorch는 여전히 데이터 복사를 수행하는 구조임

    • AMD와 같이 비-NVIDIA 디바이스에서도 ML 워크로드를 Vulkan을 통해 돌릴 수 있음

      • 최근 도입되는 cooperative matrix 같은 확장, 그리고 드라이버 레벨의 새로운 기능 지원이 추가된다면 CUDA 대비 몇 퍼센트 내외의 성능 차이도 사라질 가능성 있음

    • PyTorch의 ROCm 지원은 아주 느리게 진행 중이고, 작동에 성공한다 해도 속도가 느림

    • PyTorch가 ROCm에서 잘 작동하긴 하는데, 완전히 "동급" 수준까지 잘 되는지는 잘 모르겠음

  • PyTorch 코드에서

        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)

    대신

        self.qkv = nn.Linear(embed_size, 3*embed_size, bias = False)
    # forward 함수 내에서
    qkv = self.qkv(x)

와 같이 시도해보면 좋겠다는 제안임

  • 이렇게 하면 원래 q, k, v 각각의 파라미터가 독립적으로 연결되던 것과 달리, q, k, v 사이의 파라미터가 연결됨

    • 혹시 내가 지금 너무 피곤해서 헷갈리는 거라면 양해 바람

  • 학습자들에게 좋은 자료로 보임

    • 혹시 초보자도 따라할 수 있는 튜토리얼이나 설명서가 있는지 궁금함

    • fast.ai에서 Stable Diffusion을 직접 구현해보는 강의가 있음

  • Stable Diffusion을 라이선스 제한 없이 쓸 수 있다는 의미인지 궁금증이 생김

    • 아니고, 추론/학습 알고리즘(수학 그 자체)은 저작권 대상이 아니지만, OP는 단지 코드만 새로 구현한 것
      • 저작권이 문제되는 것은 모델(가중치) 자체고, OP는 데이터나 컴퓨팅 파워가 없어서 직접 학습하지는 않았음

  • 사실 부끄럽지만, 이저장소(repos)가 생기기 전과 후를 비교해 우리가 새롭게 얻게 된 것이 무엇인지 궁금함

    • 개인적으로 모델을 만드는 걸 피해왔고, 주로 결과물만 옆에서 지켜본 입장에서

    • 기존에도 이미 PyTorch 기반 추론/학습 스크립트가 공개되어 있는 줄로 막연히 생각했음

    • 적어도 추론 스크립트는 모델 배포시 같이 주어질 거라고 생각했고, 파인튜닝/학습 스크립트도 있을 줄 알았음

    • 이 프로젝트가 "클린룸" 혹은 "더티룸"식으로 기존 것을 다시 쓴 것인지, 아니면 기존 PyTorch 코드조차 다 CUDA/C 기반으로 너무 복잡해서 순수 PyTorch 버전이 의미가 큰 것인지 확신이 서지 않음

    • 아무튼 잘 모르겠어서 혹시 누가 설명해주면 좋겠음

      • 이 프로젝트의 핵심 가치는 "의존성이 최소화된 구현"임

        • SD 3.5를 실제로 돌려본 적은 없지만 huggingface 라이브러리를 기반으로 만들어져 있는데, 개인적으로 huggingface는 의존성이 너무 복잡해서 개발자의 환경을 거의 동일하게 맞추지 않으면 실행조차 어렵다고 느껴짐
        • 특히 오리지널 릴리즈 몇 달, 몇 년 후에는 특정 모델 실행이 매우 난이도가 높아짐
        • 예를 들어 SD3.5의 stability AI 참조구현 requirements.txt 파일 보면 버전 명시도 없고, transformers같이 엄청나게 큰 라이브러리를 포함하기 때문에 현업에서는 진짜 난감함

      • Stability AI가 Stable Diffusion 모델을 Stability AI Community License로 배포하는데, MIT와 달리 "완전 자유"는 아님

        • 특정 방식으로 가중치를 수정하는 것은 허용되지 않음
        • 이 패키지는 모델 실행(추론)이나, 혹은 이미 있는 AI 가중치로 파인튜닝은 가능한 구조임
        • 학습용/공부용으론 훌륭하지만 라이선스 이슈는 여전히 존재함

  • SD 3.5(혹은 그 어느 버전이든)를 생각할 때, 본인은 학습 과정에서 생성된 가중치(웨이트) 부분이 핵심이라고 인식함

    • 코드 자체는 결과물 품질이나 성능 측면에서 상대적으로 중요도가 떨어짐
    • 하지만 정확한 판단은 아니고, 이러한 노력을 폄하하려는 의도는 아님

  • Ludwig Maximilian University의 CompViz 그룹에서 공개한 오리지널 학술 소스의 실사용 가능성에 대해 궁금함

  • 여기 디퓨전 트랜스포머(DiT) 구현이 SD 3.5 풀 버전처럼 크로스토큰 어텐션을 제대로 구현했는지, 아니면 코드 가독성을 위해 단순화한 것인지 궁금함

답변달기