AI 이미지/ComfyUI

ComfyUI-TeaCache 사용방법

하늘이푸른오늘 2025. 5. 27. 11:29

Flux, HiDream 등 새로운 이미지 생성형 인공지능들이 예전의 모델에 비해 매개변수 수가 커지면서 생성 시간이 늘어나고 있습니다. 물론 생성되는 이미지 품질이 가장 중요하지만, 생성속도도 무시할 수 없죠.

그런데 품질에는 거의 영향을 미치지 않으면서도 디퓨전 모델의 생성 속도를 올릴 수 있는 방법이 있습니다. TeaCache를 사용하는 것입니다. TeaCache는 구현하기도 쉽고, 무엇보다 Flux와 Hunyuan Video와 같은 최신 이미지/비디오 생성모델에 적용할 수 있다는 장점이 있습니다.

제 블로그에서는 예전에도 FLUX 용 TeaCacheWan 2.1 용 TeaCache를 사용하여 생성 속도를 올리는 방법을 소개해 드린 적이 있습니다. 이 글에서는 2025년 5월말 현재 TeaCache 사용방법을 총 정리합니다.

TeaCache 란?

TeaCache는 이미지를 생성하는 동안 차(Tea)한잔을 더 마실 수 있도록 해준다는 것은 아닙니다(죄송합니다). Timestep Embedding Aware Cache(TeaCache)는 샘플링 단계에서 각 단계의 출력물간 차이를 추정하고 이를 활용하여 추론을 가속화하는, 학습이 필요없는 캐시 방식으로 디퓨전 트랜스포머(DiT, Diffusion Transformer) 모델의 샘플링 단계 속도를 올려주는 기술입니다. 이는 가장 새로운 디퓨전모델(예: Flux 및 HunyuanVideo)이 사용하고 있는 아키텍처입니다.

이제까지 생성 속도(즉 추론 속도)를 올리는 방법은 대부분 모델을 농축(distillation)하는 방법이었습니다(예: Hyper-SD, SDXL Turbo, SD 3.5 Large Turbo 등). 제가 자세한 원리는 잘 모르지만, 기본적으로는 이미지 품질을 희생하면서 샘플링 단계를 줄임으로써 속도를 올리는 방법이었습니다.

TeaCache도 "이미지 품질을 희생하면서 샘플링 단계를 줄임으로써 속도를 올리는 방법" 이라는 점에서는 농축 기법과 유사합니다. 하지만, TeaCache는 농축과는 달리 새로운 모델을 훈련시킬 필요가 없다는 것이 장점입니다. 따라서 기존에 사용하던 LoRA와 ControlNet을 모두 TeaCache와 함께 사용할 수 있습니다.

농축 모델이 빠른 이유는 품질을 상당히 희생하면서 샘플링 단계를 줄이기 때문입니다. TeaCache는 샘플링 단계는 그대로 유지하되 변화가 없는 단계를 건너뛰는 방식으로 품질 저하를 최소화합니다. cache parameter를 적절히 조절하면 눈에 띄는 품질 저하가 거의 없이 상당한 속도 향상을 이룰 수 있습니다.

품질과 속도 사이에 유연성이 있다는 것은 매우 뛰어난 특성입니다. 즉, 하나의 워크플로에서 처음에는 TeaCache를 사용해 빠르게 생성하다가 최종 결과물 생성시에만 TeaCache를 건너뛰는 방식으로 사용할 수 있습니다.

TeaCache의 작동 원리

TeaCache 는 일부 신경망 블록이 샘플링동안 많은 역할을 하지 않는다는 관찰을 활용한 것입니다. 연구자들은 디퓨전 모델이 앞쪽 샘플링 단계에서 이미지의 개략적인 윤곽을 생성하고, 나중 단계에서 세부 사항을 채운다는 것을 알아차렸습니다.

(이미지 출처: Cenn et al.)

디퓨전 트랜스포머 모델은 여러개의 트랜스포머 블록으로 구성됩니다. 앞쪽 트랜스포머 블록은 윤곽을 그리고, 뒷쪽 트랜스포머 블록은 세부사항을 그립니다. 윤곽은 나중 단계에서 변경되지 않으므로, delta-DiT 기법은 모든 단계에서 새로 계산을 하는 대신 앞쪽 블록의 생성물의 캐시를 사용하여 샘플링 단계의 속도를 올립니다.

TeaCache는 여기에서 한단계 더 나갑니다. 샘플링 동안 언제 캐시를 사용할지를 지능적으로 결정하는 것입니다. 현재의 입력이 생성된 캐시와 비슷할 경우에 생성물 캐시를 사용함으로써, 입력이 상당히 다를 때만 캐시를 새로 계산하는 것입니다. 사용자들은 얼마나 자주 캐시를 재 계산해야 할지를 제어할 수 있습니다.

TeaCache 커스톰 노드

원래의 TeaCache 깃허브에 들어가보면 FramePack 등 여러가지 프로젝트를 지원하며, 특히 ComfyUI의 경우엔 아래와 같이 여러가지 커스톰노드를 사용해 TeaCache를 사용할 수 있습니다.

이 글에서는 이중에서 가장 범용적으로 사용할 수 있는 (제일 위에 있는) ComfyUI-TeaCache 커스톰 노드에 대해 설명합니다.

ComfyUI-TeaCache 설치방법

가장 좋은 방법은 ComfyUI-Manager를 통해 설치하는 방법입니다. ComfyUI-Manager 에서 ComfyUI-TeaCache를 검색해 설치하면 됩니다.

직접 설치하시려면 아래와 같이 설치하시면 됩니다. 

  • 명령프롬프트를 띄우고 ComfyUI\custome_nodes 폴더에 들어갑니다. 
  • "git clone https://github.com/welltop-cn/ComfyUI-TeaCache.git" 명령을 실행합니다.
  • 설치가 완료된 후 ComfyUI-TeaCache 폴더에 들어갑니다.
  • "pip install -r requirements.txt" 명령을 실행합니다.

ComfyUI-TeaCache 사용방법

워크플로에 추가하는 방법

ComfyUI-TeaCache에는 현재 아래와 같이 세가지 노드가 포함되어 있습니다.

  • CompileModel
  • TeaCache
  • TeaCacheForCogVideoX

기본적인 사용 방법은 간단합니다. 그냥 Load Diffusion Model 노드 혹은 Load LoRA 노드 바로 뒤에 이들 노드를 추가해주면 됩니다. 아래는 제가 가장 많이 사용하는 Flux 용 간단 워크플로입니다.

Flux1-dev-simple.json
0.01MB

이 워크플로에서 모델을 불러오는 부분을 보면 아래와 같이 [Load Diffusion Model] -> [Load LoRA] -> [TeaCache] 순으로 연결된 것을 알 수 있습니다. 다른 모델의 경우에도 이런 방식으로 연결해 주시면 됩니다.

ComfyUI-TeaCache 지원 모델 및 매개변수 설정

TeaCache  노드의 파라미터는 아래와 같이 rel_l1_thresh, start_percent, end_percent  등이 있습니다. 

그런데, 이 커스톰 노드에서는 아래 값을 사용하라고 추천하고 있습니다. 그리고 생성된 결과의 품질이 좋지 않다고 느껴지면 rel_l1_thresh 만 낮춰보라고 권고하고 있습니다. 예를 들어 0.0을 적용하면 TeaCache를 무시하게 되므로 가장 좋은 품질을 얻을 수 있습니다. 

Models rel_l1_thresh start_percent end_percent speedup
FLUX 0.4 0 1 ~2x
PuLID-FLUX 0.4 0 1 ~1.7x
HiDream-I1-Full 0.35 0.1 1 ~2x
HunyuanVideo 0.15 0 1 ~1.9x
LTX-Video 0.06 0 1 ~1.7x
CogVideoX 0.3 0 1 ~2x
Wan2.1-T2V-1.3B 0.08 0 1 ~1.6x
Wan2.1-T2V-14B 0.2 0 1 ~1.8x
Wan2.1-I2V-480P-14B 0.26 0 1 ~1.9x
Wan2.1-I2V-720P-14B 0.25 0 1 ~1.6x
Wan2.1-T2V-1.3B-ret-mode 0.15 0.1 1 ~2.2x
Wan2.1-T2V-14B-ret-mode 0.2 0.1 1 ~2.1x
Wan2.1-I2V-480P-14B-ret-mode 0.3 0.1 1 ~2.3x
Wan2.1-I2V-720P-14B-ret-mode 0.3 0.1 1 ~2x

참고로, 여기에 들어가면 cogvideox, flux, hideram, hunyuanvideo, ltx_video, pulid_flux, wan2.1 등에 대한 예제 워크플로가 있으니 적용하실 때 참고하시면 됩니다.

Compile 모델

[Compile Model] 노드는 torch.compile을 사용하여 모델을 보다 효율적인 중간 표현(intermediate representation)으로 컴파일하여 모델 성능을 향상시킵니다. 이 컴파일 프로세스는 백엔드 컴파일러를 활용하여 최적화된 코드를 생성하므로 추론 속도를 크게 높일 수 있습니다. 처음에 워크플로를 실행할 때는 컴파일에 시간이 오래 걸릴 수 있지만, 일단 컴파일이 완료되면 추론 속도가 매우 빨라집니다.

컴파일 모델 노드를 사용하려면 아래 그림과 같이 [Load Diffusion Model] 혹은 [TeaCache] 노드 바로 뒤에 [Compile Model] 로드를 끼워주면 됩니다. 

그런데... 제가 이렇게 변경해서 실행시켜보니, 아래와 같이 W0527로 시작되는 메시지만 잔뜩 나오고(컴파일 오류인듯), 돌긴 도는데 시간 단축은 없네요. 그래서 포기하겠습니다.

이상입니다.

이 글은 ComfyUI-TeaCache 깃허브를 참고하여 작성한 글입니다.