본문 바로가기
Python

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

by teamnova 2023. 9. 15.
728x90

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

 

* 테스크 : text to image 

* 모델 : dreamshaper-v6

* 예상 결과물 테스트 방법

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

huggingface API 테스트 이미지

 

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

https://stickode.tistory.com/903

 

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

Flask 란? 파이썬으로 웹 페이지를 만들 수 있게 해주는 웹 프레임워크입니다. 보통 백엔드의 역할을 주로 수행하게 되는데요 저는 주피터 노트북에서 작성한 인공지능 모델을가지고 간단한 api

stickode.tistory.com

 

 

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)
 
    # 허깅페이스 변수
    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 가이드를 참고해주세요.