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

 

이전 글에서 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 환경에서도 유용하게 사용할 수 있는 구조라, 서비스 배포나 서버 재구동 시 매우 안정적이다.
이에 대한 근거는 차차 설명드릴 예정!