yuanchenyang/smalldiffusion

What it solves

smalldiffusion 은 확산 모델 및 흐름 모델의 학습과 샘플링 과정을 단순화하도록 설계된 경량 라이브러리입니다. 처음부터 구현하는 복잡성을 없애고, 연구자와 개발자가 새로운 모델이나 샘플러를 실험하기 쉽게 100줄 미만의 간결한 코어를 제공합니다.

How it works

라이브러리는 세 가지 주요 구성 요소의 상호 작용을 통해 동작합니다.

• Data: 학습에 PyTorch Dataset 과 DataLoader 를 사용합니다. 2D 토이 데이터셋, 픽셀 공간 이미지, 잠재 공간 이미지 데이터셋을 기본적으로 지원합니다.
• Model: MLP, U‑Net, Diffusion Transformer (DiT) 등 다양한 아키텍처를 지원하며, 스코어 예측, 흐름 예측, 데이터 예측 (x0) 등 여러 파라미터화를 선택할 수 있습니다.
• Schedule: 학습 및 샘플링 시 사용되는 노이즈 레벨(시그마)을 관리합니다. LogLinear, DDPM, LDM (Stable Diffusion 스타일), Flow matching 스케줄 등 여러 내장 스케줄을 제공합니다.

샘플링은 samples() 제너레이터 하나로 처리되며, gam 과 mu 같은 하이퍼파라미터를 조정해 DDPM, DDIM 및 기타 가속 샘플링 알고리즘을 구현할 수 있습니다.

Who it’s for

CIFAR‑10, FashionMNIST 등 소규모‑중규모 데이터셋에 대한 실험·학습이나, 사전 학습된 Hugging Face diffusers 모델에서 샘플링을 수행하고자 하는 연구자·개발자를 위한 라이브러리입니다.

Highlights

• Minimalist Core: 학습 및 샘플링 로직이 100줄 이하로 구현되어 있습니다.
• Versatile Model Support: U‑Net 및 Diffusion Transformer (DiT) 구현을 포함합니다.
• Flexible Parameterizations: 스코어, 흐름, 데이터 예측 목표를 지원합니다.
• Pretrained Model Integration: 라이브러리 샘플러를 사전 학습된 Stable Diffusion 모델과 함께 사용할 수 있는 래퍼를 제공합니다.
• High Performance: Flow matching을 사용한 ImageNet 256×256 학습 예시에서 무조건 FID가 약 27에 달하는 높은 성능을 보여줍니다.