Flask로 InstructPix2Pix 모델 서빙하기

석준·2024년 1월 19일

캡스톤디자인

목록 보기

8/9

InstructPix2Pix은 Stable Diffusion과 GPT-3를 결합한 Image-to-Image 모델로 어떠한 이미지와 명령을 입력하면 그에 맞는 이미지로 변환해주는 딥러닝 모델이다.

중장년, 노년층의 이용자가 자신의 사진을 업로드하면 얼굴을 젊게 바꿔 주거나 만화풍으로 바꿔주는 기능을 구상하였으며, 이에 위 모델이 적합하다고 생각하여 채택하였다.

캡스톤 디자인 AI 체험존, 이미지 변환 파트에 사용될 InstructPix2Pix 모델을 Flask를 사용하여 서빙 해보자!

Flask 프로젝트 생성 및 가상환경 세팅

vscode와 python 최신 버전을 준비한다.
폴더를 새로 만든다.
ctrl + shift + p를 누르고, python : create enviroment를 검색, 선택한다.
venv를 선택하고 기다린다. 해당 폴더가 생성된 것을 확인한다.
터미널을 켜서 플라스크를 다운로드한다.

python -m pip install flask

serving.py 파일을 생성하고 아래 코드를 작성한다.

from flask import Flask
app = Flask(__name__)

@app.route("/")
def home():
    return "Hello, Flask!"

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=50) 
    # port 번호 바꿔도됨
    # app.run(debug=True) -> 재시작 없이 업데이트 가능

터미널에 python serving.py를 입력하여 로컬 서버를 실행 시킨다.
http://127.0.0.1:50 에 접속해서 "Hello, Flask!" 가 잘 나온다면 성공!

InstructPix2Pix 적용

https://huggingface.co/timbrooks/instruct-pix2pix

GPU없이 CPU만으로 구동하기 위해서 Diffusers의 InstructPix2Pix을 사용했다.

pip install diffusers accelerate safetensors transformers

위와 같은 패키지가 필요하니 먼저 설치해 준다.

POST 방식으로 이미지 파일과 명령어가 전송되면 이미지를 로드한 다음, 변환을 수행하고 결과 이미지에 대한 URL 생성하여 JSON 형식으로 반환하도록 코드를 짰다.

아래는 전체 코드이다.

from flask import Flask, request, jsonify, send_from_directory
import PIL
import os
import uuid
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler

app = Flask(__name__)

model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, safety_checker=None)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

UPLOAD_FOLDER = 'uploads'
RESULT_FOLDER = 'results'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['RESULT_FOLDER'] = RESULT_FOLDER

def load_image(file_path):
    image = PIL.Image.open(file_path)
    image = PIL.ImageOps.exif_transpose(image)
    image = image.convert("RGB")
    return image

def save_image(image, folder):
    unique_filename = str(uuid.uuid4()) + '.jpg'
    file_path = os.path.join(app.config[folder], unique_filename)
    image.save(file_path)
    return file_path

@app.route('/results/<filename>')
def result_file(filename):
    return send_from_directory(app.config['RESULT_FOLDER'], filename)

@app.route('/ImageConversion', methods=['POST'])
def convert_image():
    try:
        file = request.files['file']
        prompt = request.form['prompt']

        # 업로드된 이미지를 임시로 저장
        uploaded_image_path = save_image(load_image(file), 'UPLOAD_FOLDER')
        app.logger.info(f"Uploaded file saved to: {uploaded_image_path}")

        # 이미지를 로드하고 변환 수행
        image = load_image(uploaded_image_path)
        images = pipe(prompt, image=image, num_inference_steps=8, image_guidance_scale=1.5).images
        result_image = images[0]

        # 결과 이미지를 저장
        result_image_path = save_image(result_image, 'RESULT_FOLDER')

        # 결과 이미지에 대한 URL 생성
        result_image_url = request.url_root + 'results/' + os.path.basename(result_image_path)

        return jsonify({'result': 'success', 'image_url': result_image_url})
    except Exception as e:
        return jsonify({'result': 'error', 'message': str(e)})

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=50)  # debug=True causes Restarting with stat