[Python] huggingface API로 AI 이미지생성 및 저장하기

안녕하세요 huggingface API를 통해 인공지능으로 생성한 이미지를 서버에 저장해보겠습니다. 

 

* 테스크 : text to image 

* 모델 : dreamshaper-v6

* 예상 결과물 테스트 방법

: https://huggingface.co/stablediffusionapi/dreamshaper-v6 에서 prompt를 입력해서 compute를 클릭하세요.

huggingface API 테스트 이미지

 

실행환경은 다음 글을 참고해주세요

https://stickode.tistory.com/903

[Python] Flask 를 이용해 웹 서버 실행시키기

 

 

huggingface에서 제공한 API 코드 가이드

import requests

API_URL = "https://api-inference.huggingface.co/models/stablediffusionapi/dreamshaper-v6"
headers = {"Authorization": f"Bearer {API_TOKEN}"}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.content
image_bytes = query({
	"inputs": "Astronaut riding a horse",
})
# You can access the image with PIL.Image for example
import io
from PIL import Image
image = Image.open(io.BytesIO(image_bytes))

 

 

실제 예제

@app.get('/api/letter/{prompt}')  # 라우팅할 경로를 설정해주세요

 

    print("prompt :"+prompt)

 

    # 허깅페이스 변수

    API_URL = "https://api-inference.huggingface.co/models/stablediffusionapi/dreamshaper-v6"

    headers = {"Authorization": "Bearer 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"}

    object_key_name = str(uuid.uuid4()) + '.jpg'  # 랜덤이름

    # 허깅페이스에 해당 모델에 프롬프트(text)를 전달하면 이미지를 리턴받는다.

    def query(payload):

        response = requests.post(API_URL, headers=headers, json=payload)

        return response.content

    # 프롬프트 (text)

    image_bytes = query({

        "inputs": prompt,  # 예시) "black poodle, space, ball"

    })

    # 바이트이미지 -> 이미지객체

    image = Image.open(io.BytesIO(image_bytes))

    # 이미지객체 -> 파일객체

    temp_file = '/home/ai/ai_dev_api/images/'+object_key_name+'.jpg' # 이미지를 저장할 경로 설정

    image.save(temp_file)

 

    return object_key_name 

필요한 모듈은 huggingface 가이드를 참고해주세요.