Lakonik/LakonLab

Official implementation of AsymFlow, pi-Flow, GMFlow

What it solves

LakonLab은 대규모 확산 모델을 실험하고 훈련하기 위한 고성능 프레임워크를 제공합니다. 특히 고급 플로우 매칭 기술을 구현하여 고품질 이미지를 생성하는 데 필요한 샘플링 단계를 줄여, 효율적인 몇 단계 이미지 생성 및 편집의 필요성을 해결합니다.

How it works

이 코드베이스는 세 가지 주요 연구 프로젝트를 구현합니다:

  • AsymFlow: 이미지 생성을 위한 비대칭 플로우 모델.
  • pi-Flow: 몇 단계 생성을 가능하게 하는 정책 기반 모방 증류 방법(예: FLUX 및 Qwen-Image를 위한 4단계 샘플링).
  • GMFlow: 가우시안 혼합 플로우 매칭 모델.

이를 지원하기 위해 LakonLab은 DDP, FSDP, FSDP2 간의 원활한 전환, 메모리 절약을 위한 LoRA 미세 조정 가중치 결합(weight tying), 그리고 유연한 샘플링을 위한 고급 플로우 SDE 솔버(FlowSDEScheduler 및 FlowMapSDEScheduler)를 포함합니다.

Who it’s for

대규모 확산 모델을 다루는 AI 연구자 및 개발자를 위해 설계되었습니다. 특히 플로우 매칭, 빠른 추론을 위한 증류, 그리고 고성능 훈련 인프라에 관심이 있는 분들을 대상으로 합니다.

Highlights

  • Few-Step Generation: FLUX.1 및 FLUX.2와 같은 모델에 대해 단 4단계만으로 고품질 이미지 생성을 가능하게 합니다.
  • Infrastructure Optimizations: 그래디언트 누적, 혼합 정밀도, 그리고 유연한 분산 훈련(DDP/FSDP)을 지원합니다.
  • Flexible I/O: HuggingFace, HTTP/HTTPS URLs, 및 AWS S3로부터 체크포인트를 로드하는 기능을 네이티브로 지원합니다.
  • Comprehensive Evaluation: FID, KID, IS, 및 CLIP similarity와 같은 표준 지표뿐만 아니라 GenEval 및 DPG-Bench와 같은 외부 벤치마크에 대한 통합 지원을 제공합니다.