from_pretrained 없이 로컬 구성요소만으로 Stable Diffusion 파이프라인 조립하기

root@soni 2025. 8. 7. 09:23

이전 글에서 Git LFS로 모델을 통째로 받으려다 실패한 경험을 바탕으로, 이번엔 필요한 구성요소만 뽑아서 로컬에서 파이프라인을 직접 조립해봤다.

왜 이 작업을 하게 됐을까?

처음엔 Hugging Face의 from_pretrained() 방식으로 모델을 API 호출 시마다 실시간 다운로드하는 구조를 사용했다. 간단하게 구현할 수 있는 장점은 있었지만, 다음과 같은 문제들이 생겼다.

네트워크 의존성: 컨테이너 내부나 오프라인 환경에선 사용 불가
속도 문제: 모델을 매번 로딩하느라 초기 응답이 느림
재현성 부족: 특정 시점의 모델을 안정적으로 고정해두기 어려움

그래서 모델을 미리 다운로드하고, 로컬에 저장된 구성 요소들을 직접 불러와 파이프라인을 구성하는 방식으로 전환했다.

User Request (FastAPI)
      │
      ▼
[ Load Local Components (.pth + tokenizer) ]
      │
      ▼
[ 조립된 Img2Img 파이프라인 실행 ]
      │
      ▼
[ StreamingResponse로 이미지 반환 ]

어떤 방식으로 구성했나?

Stable Diffusion 파이프라인은 다음과 같은 구성 요소들로 이루어져 있다:

UNet: 이미지 생성의 핵심
VAE: latent space ↔ 이미지 변환 담당
Text Encoder: 텍스트 프롬프트 임베딩
Tokenizer: 텍스트 토크나이징
Scheduler: 노이즈 제거 과정 스케줄링

이 중 tokenizer를 제외한 나머지는 PyTorch 모델로, .pth 파일로 저장 가능하다.

saved_sd15/
├── unet.pth
├── vae.pth
├── text_encoder.pth
└── tokenizer/
    ├── merges.txt
    ├── tokenizer_config.json
    └── vocab.json

구성 요소 저장 코드 (download_models.py)

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
)
pipe.to("cuda")

torch.save(pipe.unet, "unet.pth")
torch.save(pipe.vae, "vae.pth")
torch.save(pipe.text_encoder, "text_encoder.pth")
pipe.tokenizer.save_pretrained("tokenizer")

이렇게 하면 사전에 필요한 모든 구성 요소를 로컬 파일로 분리 저장할 수 있다.

파이프라인 직접 조립 (main.py)

unet = torch.load("unet.pth", map_location="cuda").eval()
vae = torch.load("vae.pth", map_location="cuda").eval()
text_encoder = torch.load("text_encoder.pth", map_location="cuda").eval()
tokenizer = CLIPTokenizer.from_pretrained("tokenizer")

pipe = StableDiffusionImg2ImgPipeline(
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=unet,
    scheduler=DDIMScheduler(...)
).to("cuda")

이제는 인터넷 없이도, 오직 로컬에 저장된 파일만으로 파이프라인을 조립할 수 있다.

#AI 모델 pipeline 구성

    [text prompt]                     [input image]
            │                                 │
            ▼                                 ▼
     ┌──────────────┐                  ┌──────────────┐
     │  Tokenizer   │                  │ Resize/Preproc│
     └──────┬───────┘                  └──────┬───────┘
            ▼                                 ▼
     ┌──────────────┐                  ┌──────────────┐
     │ Text Encoder │                  │     VAE      │  (Encoding)
     └──────┬───────┘                  └──────┬───────┘
            │                                 │
            ▼                                 ▼
                   ┌────────────────────┐
                   │       U-Net        │  ← 이미지 생성의 핵심
                   └────────┬───────────┘
                            ▼
                   ┌────────────────────┐
                   │     VAE (Decode)   │
                   └────────┬───────────┘
                            ▼
                      [Generated Image]

FastAPI로 감싼 후 서버를 실행하면, 네트워크 연결 없이도 이미지 생성 API가 정상 동작한다.

장점 정리

완전 오프라인 실행 가능
서버가 외부 네트워크에 연결되지 않아도 작동한다.
빠른 응답 속도
모델이 메모리에 미리 올라가 있어, 로딩 지연 없음.
환경 재현성
특정 시점의 모델 구성요소를 고정해 둘 수 있어 안정적인 배포 가능.

다음에 개선할 점

현재 구성 요소는 torch.save()로 저장했기 때문에, PyTorch와 diffusers의 버전이 달라지면 ModuleNotFoundError가 발생할 수 있다. (나에게 이 이슈가 발생해서 다음 게시글은 Error 핸들링에 관한 내용입니다..)
안정성을 더 높이려면 구성 요소를 state_dict 형태로 저장하는 방식도 고려해볼 수 있다 (이건 나중에 따로 다룰 예정).

마무리

이번 방식은 Docker 기반 MLOps 환경에서도 유용하게 사용할 수 있는 구조라, 서비스 배포나 서버 재구동 시 매우 안정적이다.
이에 대한 근거는 차차 설명드릴 예정!