3-2학기, 졸업 프로젝트 스타트를 잘 마치고 여러 회의와 서베이를 거쳐 졸업 프로젝트의 주제를 변경하기로 했다.
이전 주제는 스키밍 장치 탐지를 위한 어플리케이션을 만드는 것으로 이상탐지 모델을 사용하여 구현하려 하였지만, 스키밍 장치에 대한 데이터셋 부족과 스키밍 장치마다 모델을 만들기는 어렵다는 판단이 들어 새로운 프로젝트를 진행하게 되었다.
우리의 새로운 프로젝트는 바로
전시 관람 후 사진과 관람평을 기록하고 나만의 티켓 굿즈를 만들어주는 아카이빙 웹 서비스
이다!
이 중, 중요 기능 중 하나인 전시 스타일을 반영하여 나만의 티켓 굿즈 생성하는 기능에서 가장 주요한 부분인 티켓 이미지를 생성하는 부분에 대해 기술하려고 한다.
정확한 목표는 사용자가 찍은 전시회 사진을 입력받아, 해당 전시 스타일을 반영한 티켓(에 들어갈) 이미지를 생성하는 것!
Step1. AI 모델 선정하기
스타일을 반영한 이미지를 생성하기 위해선 style transfer AI 모델을 사용해야하기 때문에 최신 style transfer 모델을 서베이해보았다.
그러다 'Styleshot'이라는 모델을 알게 되었는데
해당 모델은 두 개의 input 이미지와 prompt를 통해 새로운 이미지를 생성한다.
input 이미지는 style reference와 content reference로 나뉘어져 있다.
- style reference -> 스타일을 추출
- content reference -> 윤곽을 추출
이렇게 두 가지의 이미지에서 각기 다른 요소를 추출하여 합쳐주는 것이 이 styleshot이 작동하는 방식인 것이다.
test 해본 결과
전시회의 스타일을 잘 추출하여 자연스럽게 만들어주는 것이, 우리 프로젝트에서 사용하기에 적합하다고 생각이 들어 모델을 사용할 준비를 하였다.
Step2. local에서 AI 모델 사용하기/구축하기
먼저 모델을 다운받아 local에서 잘 돌아가는지 확인하는 작업을 진행하였다.
git clone
git clone https://github.com/Jeoyal/StyleShot.git
cd StyleShot환경설정
conda create -n styleshot python==3.8
conda activate styleshot
pip install -r requirements.txtpre-trained weight 다운로드
git lfs install
git clone https://huggingface.co/Gaojunyao/StyleShot_lineart모델을 inference하기 위해선 pre-trained weight를 따로 다운받아야한다.
styleshot에는 lineart 와 contour 2가지의 모드가 있는데, 이중 우리는 lineart 모드를 사용하기 때문에 해당 모드의 weight만 다운 받았다.
-> 오픈 소스 모델이므로 github에 잘 정리가 되어있어 클론과 환경 설정까지 진행을 한 후 inference하는 파일을 실행해보았다.
하지만,, 에러가 났고 해당 에러는 huggingface에서의 버전이 달라져 생기는 오류였다.
에러가 생기는 파일 수정
- ./ip_adapter/ip_adapter.py
해당 파일에서 오류가 되는 부분을 수정해준 뒤, 다시 inference하는 파일을 실행해본 결과 잘 실행이 되었다.
-
style reference image
-
content reference image
-
output image
우리가 원하던 대로 전시 스타일을 잘 반영하면서 티켓 이미지에 사용할 사진이 잘 생성이 되었다.
step3. AI 모델 배포하기
local 환경에서 모델이 잘 작동하는 것을 확인했으니, 이제 우리의 웹 서비스에서 사용할 수 있도록 모델을 배포해야한다.
AI 모델을 배포하는 방법은 여러가지가 있는데, 우리 sw구조와 알맞은 방법으로 FastAPI를 통해 AI 서버로 모델을 배포하는 방식을 사용하였다.
memoir의 루트 폴더를 만들어준 뒤,
그 안에 방금 우리가 Styleshot 모델을 사용할 때 다운받아두었던 Styleshot 폴더와 FastAPI로 배포하는 main.py 파일을 넣어준다.
- main.py
import sys
import os
import requests
sys.path.append(os.path.join(os.path.dirname(__file__), "StyleShot"))
import io
import torch
import cv2
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse
from PIL import Image
from StyleShot.annotator.hed import SOFT_HEDdetector
from StyleShot.annotator.lineart import LineartDetector
from diffusers import UNet2DConditionModel, ControlNetModel
from transformers import CLIPVisionModelWithProjection
from huggingface_hub import snapshot_download
from StyleShot.ip_adapter import StyleShot, StyleContentStableDiffusionControlNetPipeline
import boto3
from botocore.exceptions import NoCredentialsError
import uuid
from PIL import Image
from typing import Optional
from dotenv import load_dotenv
load_dotenv()
def upload_to_s3(local_file_path: str, bucket: str, region: str, s3_key_prefix: str = "") -> str:
s3 = boto3.client('s3')
filename = f"{s3_key_prefix}{uuid.uuid4().hex}.png"
try:
s3.upload_file(local_file_path, bucket, filename)
return f"https://{bucket}.s3.{region}.amazonaws.com/{filename}"
except NoCredentialsError:
raise RuntimeError("S3 credentials not found.")
app = FastAPI()
# === 모델 초기화 ===
base_model_path = "models/stable-diffusion-v1-5"
transformer_block_path = "models/CLIP-ViT-H-14-laion2B-s32B-b79K"
device = "cuda" if torch.cuda.is_available() else "cpu"
styleshot_models = {}
def load_styleshot(preprocessor: str):
if preprocessor in styleshot_models:
return styleshot_models[preprocessor]
if preprocessor == "Lineart":
detector = LineartDetector()
styleshot_model_path = "models/StyleShot_lineart"
elif preprocessor == "Contour":
detector = SOFT_HEDdetector()
styleshot_model_path = "Gaojunyao/StyleShot"
else:
raise ValueError("Invalid preprocessor")
# 모델 다운로드
if not os.path.isdir(styleshot_model_path):
snapshot_download(styleshot_model_path, local_dir=styleshot_model_path)
if not os.path.isdir(base_model_path):
snapshot_download(base_model_path, local_dir=base_model_path)
if not os.path.isdir(transformer_block_path):
snapshot_download(transformer_block_path, local_dir=transformer_block_path)
ip_ckpt = os.path.join(styleshot_model_path, "pretrained_weight/ip.bin")
style_aware_encoder_path = os.path.join(styleshot_model_path, "pretrained_weight/style_aware_encoder.bin")
unet = UNet2DConditionModel.from_pretrained(base_model_path, subfolder="unet")
content_fusion_encoder = ControlNetModel.from_unet(unet)
pipe = StyleContentStableDiffusionControlNetPipeline.from_pretrained(
base_model_path,
controlnet=content_fusion_encoder
)
styleshot = StyleShot(device, pipe, ip_ckpt, style_aware_encoder_path, transformer_block_path)
styleshot_models[preprocessor] = (styleshot, detector)
return styleshot, detector
# === API 엔드포인트 ===
@app.post("/generate/")
async def generate_image(
preprocessor: str = Form(...),
style_url: str = Form(...),
content_url: str = Form(...),
prompt: Optional[str] = Form(None)
):
if prompt is None:
prompt = "default prompt" # 또는 그냥 빈 문자열로 처리
styleshot, detector = load_styleshot(preprocessor)
# 스타일 이미지 다운로드
style_response = requests.get(style_url)
style_response.raise_for_status()
style_image = Image.open(io.BytesIO(style_response.content)).convert("RGB")
# 콘텐츠 이미지 다운로드
content_response = requests.get(content_url)
content_response.raise_for_status()
content_array = np.frombuffer(content_response.content, np.uint8)
content_image = cv2.imdecode(content_array, cv2.IMREAD_COLOR)
content_image = cv2.cvtColor(content_image, cv2.COLOR_BGR2RGB)
processed_content = detector(content_image)
processed_content = Image.fromarray(processed_content)
# 추론
result = styleshot.generate(style_image=style_image, content_image=processed_content)
output_image = result[0][0]
# 결과 반환
img_bytes = io.BytesIO()
output_image.save(img_bytes, format='PNG')
img_bytes.seek(0)
output_path = f"/tmp/{uuid.uuid4().hex}.png"
output_image.save(output_path) # 파일 저장
# S3 업로드
s3_url = upload_to_s3(
local_file_path=output_path,
bucket="hukmemoirbucket",
region="ap-northeast-2",
s3_key_prefix="results/"
)
os.remove(output_path) # 서버 공간 정리
return JSONResponse(content={"s3_url": s3_url})
@app.post("/test-upload/")
def test_s3_upload():
# 1. 더미 이미지 생성
dummy_image = Image.new("RGB", (256, 256), color="blue")
temp_path = "test_result.png"
dummy_image.save(temp_path)
# 2. S3 업로드
try:
s3_url = upload_to_s3(
local_file_path=temp_path,
bucket="hukmemoirbucket",
region="ap-northeast-2",
s3_key_prefix="test/"
)
os.remove(temp_path) # 업로드 후 파일 정리
return {"s3_url": s3_url}
except Exception as e:
return {"error": str(e)}
이 main.py 파일에서는 모델을 FastAPI로 배포해주고, 이미지를 받아 티켓 이미지를 생성하게 되면 이를 S3에 저장하고 해당 url를 반환해주는 코드까지 구현되어있다.
이때 local에서와 다른 점은 이 memoir를 local에서 돌렸던 상태로 git에 올리게 되면 너무나도 큰 용량 때문에 올라가지 않게 된다.
이를 해결하기 위해 아까 다운받은 pre-trained weight 파일들은 .gitignore 처리를 한 후
ec2 서버에서 위의 모델들을 다시 다운 받을 수 있게 하는 코드를 작성해둔다.
- download_weights.sh
#!/bin/bash
set -e # 에러 발생 시 중단
echo "🔽 StyleShot 모델 및 종속 모델 다운로드 시작..."
# 기본 디렉토리 구조 설정
MODEL_DIR="models"
mkdir -p $MODEL_DIR
# huggingface-cli 설치 확인
if ! command -v huggingface-cli &> /dev/null; then
echo " huggingface-cli가 설치되어 있지 않습니다. 설치 중..."
pip install -q huggingface_hub
fi
# git-lfs 설치 확인
if ! command -v git-lfs &> /dev/null; then
echo " git-lfs가 설치되어 있지 않습니다. 설치 중..."
sudo apt-get update && sudo apt-get install -y git-lfs
git lfs install
fi
# 모델 리스트
declare -A models=(
["StyleShot_lineart"]="Gaojunyao/StyleShot_lineart"
["StyleShot_contour"]="Gaojunyao/StyleShot"
["stable-diffusion-v1-5"]="runwayml/stable-diffusion-v1-5"
["CLIP-ViT-H-14-laion2B-s32B-b79K"]="laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
)
# 모델 다운로드
for dir in "${!models[@]}"; do
path="$MODEL_DIR/$dir"
repo="${models[$dir]}"
if [ ! -d "$path" ]; then
echo " $repo → $path"
python3 -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='$repo', local_dir='$path', local_dir_use_symlinks=False)"
else
echo " $repo 이미 존재: $path"
fi
done
echo " 모든 모델 다운로드 완료!"
이렇게 해두면 git에는 base 모델이 올라가있지 않지만, ec2 서버에서 git clone을 받은 수 download_weights.sh을 실행시키면 자동으로 models 폴더가 만들어지고 여기에 필요한 모델들이 다운로드 된다.
이렇게 처음 모델을 다운시켜두고 한번 inference를 하게 되면 그 다음부터는 checkpoint가 생겨 빠르게 이미지를 생성할 수 있게 된다.
ec2 서버에서 실행해보기
이제 github에 올린 memoir 폴더를 ec2 서버에서 클론 받은 뒤
FastAPI로 배포한 styleshot을 실제 ec2 서버에서 실행시켜보았을 때
잘 실행이 되었고 생성된 이미지가 S3에 저장이 되고 해당 url를 return 하는 것까지 잘 실행되는 것이 확인되었다
이후 이를 docker로 묶어 배포하는 것은 다른 팀원의 역할이였기에 이 상태로 넘겨주며 프로젝트를 마치게 되었다.
작성자 : (Team 28 HUK) 2171008 김지수
