lucidrains/rectified-flow-pytorch

해결하는 문제

이 프로젝트는 노이즈와 데이터 사이의 경로를 더 직선적으로 만들어 샘플링 과정을 단순화하고 효율성을 높이기 위해 설계된 생성 모델링 기법인 Rectified Flow의 PyTorch 구현을 제공합니다.

작동 방식

핵심 Rectified Flow 알고리즘과 이후 연구에서 제안된 개선 사항들을 구현합니다. 사용자는 모델(예: UNet)을 정의하고 RectifiedFlow 클래스로 감싸면 학습 손실과 샘플링을 처리할 수 있습니다. 또한 "reflow"라는 반복적으로 흐름을 더 직선화하는 과정을 지원하여, 더 빠른 생성이 가능해질 수 있습니다.

대상 사용자

이미지 합성에서 플로우 매칭 및 전통적인 확산 모델의 대안에 관심이 있는 생성형 AI 연구자 및 개발자.

주요 특징

• Reflow 지원: 경로를 반복적으로 직선화하는 Reflow 클래스를 포함합니다.
• 통합 트레이너: accelerate 라이브러리 기반의 Trainer를 제공하여 모델 학습을 더 쉽게 합니다.
• 다양한 구현: LSD 플로우, SoFlow, Split Mean Flow, Unconstrained Alignment (UA) 플로우, 라플라시안 다중스케일 플로우 매칭 등의 고급 변종을 포함합니다.
• 즉시 사용 가능한 컴포넌트: 빠른 실험을 위해 내장된 Unet과 ImageDataset를 제공합니다.