Stable Diffusion 이미지 생성 API 만들기
회사에서 도입한 GPU 머신에 구축한 환경 위에서, Stable Diffusion 모델을 FastAPI로 감싸 간단한 이미지 생성 API를 구현했다.
이미지 기반 AI 기능을 실서비스처럼 테스트하려면, 단순히 모델을 로딩하는 것만으론 부족하다. API 형태로 만들어야 실제로 쓸 수 있기 때문에, FastAPI를 활용해 Stable Diffusion 모델을 쉽게 테스트하고 실행할 수 있도록 구성했다.
Hugging Face 연동
모델을 허브에서 다운로드하기 위해선 CLI 인증이 필요하다.
pip install huggingface_hub
huggingface-cli login
토큰 발급 후 로그인하면 끝. (토큰은 huggingface 사이트에 접속해서 쉽게 발급 가능하다)
Stable Diffusion 파이프라인 구성
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16
).to("cuda")
이번 구현에서는 Hugging Face의 runwayml/stable-diffusion-v1-5 모델을 기반으로, img2img 파이프라인을 사용했다. 업로드된 이미지를 입력으로 받아 프롬프트에 맞춰 변환하는 방식으로 작동하며, GPU에서 빠르게 추론이 가능하도록 float16 설정과 함께 .to("cuda")로 올려준다.
현재는 from_pretrained()를 통해 모델과 관련 구성 요소들을 매번 Hugging Face Hub에서 다운로드하고 있지만, 추후에는 모델 구성 요소들(unet, vae, tokenizer 등)을 미리 로컬에 저장한 뒤 불러오는 방식으로 전환할 계획이다.
이렇게 하면 네트워크 의존성을 줄이고, 컨테이너 환경에서 보다 안정적으로 실행할 수 있다는 장점이 있다. 해당 방식은 다음 글에서 구체적으로 다룰 예정이다.
FastAPI 엔드포인트 작성
@app.post("/generate")
async def generate(prompt: str = Form(...), image: UploadFile = File(None)):
...multipart/form-data로 프롬프트와 이미지를 함께 받아서, 변환된 이미지를 StreamingResponse로 바로 응답해준다.
Swagger UI로 바로 테스트 가능
uvicorn ExampleFile:app --reload
실행 후 브라우저에서 http://localhost:8000/docs 로 들어가면 API를 바로 시각화해서 테스트할 수 있다.
재실행 시 요약 커맨드
conda activate photobooth uvicorn ExampleFile:app --reload정리하자면
- 현재는 애플리케이션 실행할 때마다 모델과 관련 구성 요소들을 Hugging Face Hub에서 다운로드하여 사용하고 있지만 다음 게시글에서는 모델 구성 요소들(unet, vae, tokenizer 등)을 미리 로컬에 저장한 뒤 불러오는 방식으로 전환할 것이다.
- Hugging Face 모델을 로컬 GPU에서 빠르게 써볼 수 있는 API 형태로 구성했다.
- Swagger UI를 통해 시각적으로 테스트할 수 있어서 디버깅이 편하다.
- 이렇게 AI 이미지 처리 라이브러리들은 쉽게 테스트 가능하고, gpt Chat으로 요청하는 것보다 퀄리티는 떨어지지만 훨씬 재밌다!