AudioLDM 2 최적화 가이드: Hugging Face Diffusers를 사용한 추론 시간 단축

Hugging Face는 AudioLDM 2에 대한 최적화 가이드를 출시했습니다. AudioLDM 2는 현실적인 사운드 효과, 인간 음성, 음악을 생성할 수 있는 텍스트-오디오 잠재 확산 모델입니다. diffusers 라이브러리 내에서 코드와 모델 최적화를 조합하여 적용함으로써, 10초 오디오 샘플의 추론 시간이 원래 구현에서의 30초 이상에서 1초 미만으로 줄어듭니다.

AudioLDM 2 모델 아키텍처

AudioLDM 2는 텍스트 임베딩으로부터 연속적인 표현을 학습하여 오디오를 생성하는 텍스트-오디오 잠재 확산 모델(LDM)입니다. 생성 과정은 다단계 파이프라인을 따릅니다:

  1. 텍스트 인코딩: 모델은 CLAP의 텍스트 브랜치(오디오 정렬 임베딩용)와 Flan-T5의 텍스트 인코더(의미적 표현용) 두 개의 인코더를 사용하여 텍스트 임베딩을 계산합니다.
  2. 프로젝션: 이러한 임베딩은 AudioLDM2ProjectionModel을 통해 선형 투영을 거쳐 공유 임베딩 공간으로 전달됩니다.
  3. 자기회귀 생성: GPT2 언어 모델은 투영된 CLAP과 Flan-T5 임베딩에 조건부인 $N$개의 임베딩 벡터 시퀀스를 생성합니다.
  4. 잠재 확산: UNet 기반 LDM은 $T$ 추론 단계 동안 랜덤 잠재를 디노이즈합니다. 대부분의 LDM과 달리, AudioLDM 2 UNet은 GPT2와 Flan-T5에서 나온 두 세트의 크로스 어텐션 임베딩을 사용합니다.
  5. 디코딩: 최종 디노이즈된 잠재는 VAE 디코더를 통해 멜 스펙트로그램을 복구한 후, 보코더를 통해 오디오 파형으로 변환됩니다.

사용 가능한 모델 체크포인트

체크포인트 태스크 모델 크기 훈련 데이터 (시간)
cvssp/audioldm2 텍스트-오디오 1.1B 1150k
cvssp/audioldm2-music 텍스트-음악 1.1B 665k
cvssp/audioldm2-large 텍스트-오디오 1.5B 1150k

속도를 위한 추론 최적화

원래 구현의 느린 추론 속도를 해결하기 위해, Hugging Face는 diffusers 라이브러리에서 사용할 수 있는 네 가지 주요 최적화 기술을 식별합니다:

1. 플래시 어텐션 (SDPA)

PyTorch 2.0 이상을 사용하면, diffusers 라이브러리가 자동으로 torch.nn.functional.scaled_dot_product_attention (SDPA)를 활성화합니다. 이는 플래시 어텐션과 유사한 메모리 효율적인 어텐션 연산을 제공하여 출력 품질을 변경하지 않고 계산 시간을 줄입니다.

2. 반정밀도 (float16)

모델 가중치와 연산을 float32에서 float16(반정밀도)으로 변환하면 GPU 메모리 사용량이 크게 줄어들고 추론 속도가 향상되며, 음질에 거의 영향을 미치지 않습니다. 이는 .from_pretrained 호출 중에 torch_dtype=torch.float16을 전달함으로써 달성됩니다.

3. 토치 컴파일

파이프라인에서 가장 계산 비용이 높은 UNet 부분을 torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)로 감싸면 상당한 속도 향상을 얻을 수 있습니다. 첫 번째 추론 실행은 컴파일 오버헤드로 인해 느릴 수 있습니다(잠재적으로 2분까지), 하지만 이후 모든 세대는 znacz히 빨라집니다.

4. 효율적인 스케줄러

기본 DDIMScheduler(일반적으로 200단계가 필요)를 DPMSolverMultistepScheduler와 같은 더 성능 좋은 스케줄러로 교체하면 모델은 20-25개의 추론 단계만으로 유사한 품질을 달성할 수 있습니다. 이 최적화를 다른 최적화들과 결합하면 10초 샘플의 생성 시간을 1초 미만으로 줄일 수 있습니다.

메모리 관리 및 장형 오디오

150초(예시)와 같은 긴 오디오 샘플을 생성하거나 audioldm2-large와 같은 더 큰 체크포인트를 사용하면 잠재 변수의 폭이 증가하여, 크로스 어텐션 메모리가 시퀀스 길이에 따라 2차적으로 증가하기 때문에 CUDA 메모리 부족(OOM) 오류가 발생할 수 있습니다.

이를 완화하기 위해, Hugging Face는 pipe.enable_model_cpu_offload()를 통한 CPU 오프로딩을 권장합니다. 이 기술은 현재 활성화된 모델 구성 요소만 GPU에 유지하고 나머지를 CPU로 오프로드하여, 제한된 RAM을 가진 GPU에서도 장형 오디오 생성과 더 큰 모델 사용을 가능하게 하며, 추론 시간에 최소한의 페널티만 발생시킵니다.

Sources