Stability AI API의 Inpaint 기능으로 그림의 일부 파트만 생성하기

 

 

뛰어난 이미지 생성 성능으로 잘 알려진 diffusion 기반 이미지 생성 모델 stable diffusion은 오픈 소스로 공개되어 있기도 하지만, 모델을 운용할 인프라가 부족하다면 Stability AI 사의 API 서비스로 이용할 수 있기도 합니다.
 
Stability AI API 서비스 페이지
 
API 서비스는 비용이 부과되지만, 간편하게 앱을 구축할 수 있다는 장점이 있습니다. 비용은 크레딧(credit)을 충전하는 형태이며, 10$에 1,000 크레딧을 구매할 수 있습니다. Pricing 메뉴에서 task 별 소모되는 크레딧을 볼 수 있는데, 업스케일이나 영상 생성 같은 무거운 작업이 아니라면 한 번에 4-8 크레딧 정도가 소모되는 것 같습니다. 1,000 크레딧 정도면 개인 프로젝트용으로는 충분합니다. 카드를 등록하여 크레딧을 구매하면 됩니다. 본 포스트에서는 Stability AI API를 이용해 여러 기능 중 inpaint 생성 기능을 테스트해보겠습니다.
 

Inpaint 란?

 
Inpaint는 그림을 일부 영역만 지정하여 해당 영역에만 이미지 생성을 수행하는 기능입니다. 생성할 영역 지정은 보통 흑백으로 이루어진 마스크 이미지(mask image)를 이용합니다. 모델은 마스크 이미지의 흰색 영역만을 인식하여 생성을 수행하게 됩니다. 물론, 원본 이미지와 마스크 이미지의 크기는 같아야겠습니다.
 

출처: AWS blog[1]

 

마스크 이미지 만들기

 
마스크 이미지는 정확한 경계가 필요하지 않기 때문에, 간단히 만들 수 있습니다. 먼저 inpaint를 수행할 원본 이미지를 준비합니다.

 
그림판이나 간단한 편집도구로 생성에서 먼저 이미지를 검은색으로 칠한 뒤, 생성할 영역을 흰색으로 덧칠해줍니다. 이미지 생성 모델이 주변 경계와 어느 정도 어울리게 생성을 수행하므로, 경계를 여유있게 설정해도 괜찮습니다. 위 이미지의 마스크 이미지 예제는 다음과 같습니다.

 
 

Stability AI API 호출

 
이제 코드를 작성해 inpaint를 수행해봅시다. requests 라이브러리로 간단히 API를 이용할 수 있습니다. API 호출은 공식 문서를 참고하면 됩니다.
 
Stability AI API Reference 페이지
 

import requests
from IPython.core.display import Image, display

api_key = '사용자 API 키 입력'
output_format = 'png' # 출력 이미지 확장자 지정

 
원본 이미지와 마스크 이미지를 확인해봅시다.
 

display(Image(filename = "cake_template.png", width=200, unconfined=True))
display(Image(filename = "cake_template_mask.png", width=200, unconfined=True))

 
프롬프트는 다음과 같이 설정해봅시다. Stability AI의 모델은 한글 입력을 지원하지 않으므로, 영어로 작성해줍시다.
 

prompt = "designed cake, blue waves and tides, dolphin, 
'Happy Birthday Jane' message, some star-shaped decorations."

 
 
이제 API를 호출하여 이미지를 생성해봅시다. 생성된 이미지를 현재 경로에 저장합니다. 
 

response = requests.post(
    f"https://api.stability.ai/v2beta/stable-image/edit/inpaint",
    headers={
        "authorization": f"Bearer {api_key}",
        "accept": "image/*"
    },
    files={
        "image": open("cake_template.png", "rb"),
        "mask": open("cake_template_mask.png", "rb"),
    },
    data={
        "prompt": prompt,
        "output_format": output_format,
    },
)

if response.status_code == 200:
    with open(f"cake_inpainted.{output_format}", 'wb') as file:
        file.write(response.content)
else:
    raise Exception(str(response.json()))

 
 
생성된 이미지를 확인해봅시다.
 

display(Image(filename = "cake_inpainted.png", width=600, unconfined=True))

 
배경을 제외하고, 마스크 이미지에서 흰색 영역만 프롬프트대로 잘 생성된 것 같습니다. 또한, 마스크된 영역과 아닌 영역 사이의 경계도 어느 정도 잘 보정되어 생성이 됩니다. Inpaint 기능을 이용하면 원본 이미지에서 살리고 싶은 부분을 보존하면서 이미지 생성을 수행할 수도 있고, 복수의 이미지 생성 시 일관성 있는 템플릿 유지가 가능한 장점이 있습니다.
 
 
[1] https://aws.amazon.com/ko/blogs/machine-learning/inpaint-images-with-stable-diffusion-using-amazon-sagemaker-jumpstart/