Diffusers와 함께 무료 티어 Google Colab에서 DeepFloyd IF 실행하기

TL;DR

DeepFloyd의 IF 모델은 T5‑XXL 텍스트 인코더를 8‑비트로 양자화하고 🤗 Diffusers로 모델 구성 요소를 모듈식으로 로드하며 가중치를 CPU 또는 디스크로 오프로드함으로써, 노트북의 13 GB RAM 및 15 GB GPU 제한에도 불구하고 전체 텍스트‑투‑이미지, 이미지‑변형, 인페인팅 파이프라인을 실행할 수 있습니다.

소개 – IF의 기능과 제약

DeepFloyd는 2023년 4월 말에 IF를 공개했으며, 이는 Google의 Imagen에서 영감을 받은 픽셀‑스페이스 텍스트‑투‑이미지 확산 모델입니다. Stable Diffusion과 비교했을 때 IF는:

압축되지 않은 이미지에 직접 작동하여 얼굴·손과 같은 고주파 디테일을 보존합니다.
CLIP 대신 강력한 T5‑XXL 인코더를 사용해 텍스트 충실도를 높이며, 읽을 수 있는 텍스트를 안정적으로 생성하는 최초의 오픈‑소스 모델이 되었습니다.

그 대가로 모델 크기가 큽니다: T5‑XXL(45억 파라미터), IF‑I UNet(43억), IF‑II 업스케일러 UNet(12억) 모두 합쳐 100억 파라미터를 초과하며, Stable Diffusion 2.1의 약 13억 파라미터와는 큰 차이가 있습니다. 전체 모델을 float32로 실행하려면 40 GB 이상의 GPU 메모리가 필요하지만, 무료 Colab T4(15 GB VRAM)와 13 GB CPU RAM에서는 이를 제공할 수 없습니다.

메모리 제한 하드웨어를 위한 최적화

핵심 도구

🤗 Accelerate – 대형 모델의 디바이스 배치를 위한 유틸리티.
bitsandbytes – PyTorch 모델을 위한 8‑비트 양자화.
🤗 safetensors – 빠르고 안전한 체크포인트 로딩.
🤗 Diffusers – 위 도구들을 통합한 고수준 확산 파이프라인.

메모리 절감 전략

T5‑XXL을 8‑비트로 양자화 – 인코더 체크포인트를 ~20 GB(fp32)에서 ~8 GB(8‑bit safetensors)로 감소.
구성 요소를 지연 로드 – Diffusers는 텍스트 인코더만 혹은 UNet만 로드하도록 허용해 동시에 메모리 피크가 발생하지 않게 함.
UNet 가중치를 fp16으로 사용 – Stage 1·Stage 2 UNet을 반정밀도로 로드하면 GPU 메모리에 들어감.
PyTorch 객체를 명시적으로 해제 – 각 단계 후 del 로 모델 객체를 삭제하고 gc.collect()와 torch.cuda.empty_cache() 를 실행.

단계별 텍스트‑투‑이미지 생성

1. 최신 의존성 설치

pip install --upgrade \
  diffusers~=0.16 \
  transformers~=4.28 \
  safetensors~=0.3 \
  sentencepiece~=0.1 \
  accelerate~=0.18 \
  bitsandbytes~=0.38 \
  torch~=2.0 -q

2. 8‑비트 T5 인코더 로드

from transformers import T5EncoderModel
text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder",
    device_map="auto",
    load_in_8bit=True,
    variant="8bit",
)

3. UNet을 로드하지 않고 프롬프트 임베딩 생성

from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=text_encoder,
    unet=None,
    device_map="auto",
)
prompt = "a photograph of an astronaut riding a horse holding a sign that says Pixel's in space"
prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

4. UNet을 위한 메모리 확보를 위해 인코더 해제

import gc, torch

def flush():
    gc.collect()
    torch.cuda.empty_cache()

del text_encoder, pipe
flush()

5. Stage 1 확산 (64×64) – UNet만 로드

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=None,
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="auto",
)
generator = torch.Generator().manual_seed(1)
image = pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    output_type="pt",
    generator=generator,
).images

결과는 pt_to_pil 로 시각화할 수 있는 64 × 64 텐서입니다.

6. Stage 2 초고해상도 (64→256) – IF‑II 파이프라인

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0",
    text_encoder=None,
    variant="fp16",
    torch_dtype=torch.float16,
    device_map="auto",
)
image = pipe(
    image=image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds,
    output_type="pt",
    generator=generator,
).images

7. Stage 3 초고해상도 (256→1024) – StabilityAI x4 업스케일러

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler",
    torch_dtype=torch.float16,
    device_map="auto",
)
final_image = pipe(prompt, image=image, generator=generator).images[0]

원한다면 IF 워터마크를 수동으로 적용:

from diffusers.pipelines.deepfloyd_if import IFWatermarker
watermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")
watermarker.apply_watermark(final_image, pipe.unet.config.sample_size)

이제 파이프라인은 무료 Colab 세션 내에서 1024 × 1024 이미지를 완전히 생성합니다.

이미지 변형 및 인페인팅 – 동일 체크포인트 재사용

IF 체크포인트는 IFImg2ImgPipeline(변형)과 IFInpaintingPipeline도 지원합니다. 워크플로는 텍스트‑투‑이미지 단계와 동일합니다:

8‑비트 T5 인코더를 로드하고 변형 프롬프트를 인코딩한 뒤 인코더를 해제.
변형 파이프라인을 위해 unet=None 로 Stage 1 UNet을 로드하고 원본 이미지와 strength 를 전달해 노이즈 양을 조절.
IFImg2ImgSuperResolutionPipeline(또는 Stable Diffusion 업스케일러)으로 업스케일.
인페인팅의 경우 추가로 바이너리 마스크 이미지를 제공; 동일한 모듈식 로딩·메모리 해제 패턴을 적용.

세 파이프라인 모두 동일한 메모리 최적화 트릭을 공유합니다: 8‑비트 텍스트 인코더, fp16 UNet, 자동 디바이스 매핑, 명시적 가비지 컬렉션.

실용적인 고려사항 및 성능 트레이드‑오프

속도 vs. 메모리 – 8‑비트 양자화와 반복적인 로드/언로드는 추론 지연을 증가시킵니다. 프로덕션 환경에서는 ≥40 GB VRAM을 갖춘 GPU(예: A100)를 사용해 모든 구성 요소를 디바이스에 유지하면 최대 처리량을 확보할 수 있습니다.
품질 – Hugging Face Spaces에 호스팅된 공식 IF 데모는 전체 정밀도 모델을 실행하므로 특히 대형 이미지에서 약간 더 높은 충실도를 제공합니다.
라이선스 – 체크포인트를 로드하기 전에 모델 카드에서 DeepFloyd IF 라이선스를 반드시 수락해야 합니다.

결론 – 대형 확산 모델의 민주화

오픈‑소스 체크포인트(DeepFloyd IF, StabilityAI 업스케일러)와 커뮤니티 주도 라이브러리(Diffusers, Transformers, Accelerate, bitsandbytes)를 결합함으로써 Hugging Face는 10 B‑파라미터 이상 모델을 무료 Google Colab 인스턴스에서도 실행할 수 있음을 보여줍니다. 이는 모듈식 파이프라인과 양자화가 최첨단 생성 AI에 대한 접근성을 넓히는 강력한 수단임을 입증합니다.

Diffusers와 함께 무료 티어 Google Colab에서 DeepFloyd IF 실행하기