Streamlit ? AI/ML/DL Demo 사용시 유용 Tool

AI 개발을 하면서 미리 Demo Version을 만들어야 하는 경우들이 종종 있다. 이럴때 간편하고 빠르게 AI Demo 를 공유해서 서 Server까지 열어 줄수 있는 데 이 때 유용한 Streamlit을 활용해 보는 예시를 들어보겠다.

일단 streamlit은 주로 데이터시각화에 유용하게 맞춰줘있고 다양한 database를 화면에 바로바로 보여줄수있는 편리한 기능까지 마련되어있다

 

• 공식사이트는 아래 link에서 확인하세용
  • https://docs.streamlit.io/

Streamlit Docs

 

 

• Streamlit install

pip install streamlit
streamlit hello

 

• Streamlit base

import streamlit as st

• install 한다음 strealit 모듈을 활용해 간편하게 여러가지 화면 렌더링이 가능하다
• rendering은 st.write() 을 사용해 거의 모든 것들이 가능한 수준이다
• 예를 들어서
•  

import streamlit as st
import pandas as pd
import numpy as np 

st.write("Here's our first attempt at using data to create a table:")
st.write(pd.DataFrame({
    'first column': [1, 2, 3, 4],
    'second column': [10, 20, 30, 40]
}))

dataframe = np.random.randn(10, 20)
st.dataframe(dataframe)

chart_data = pd.DataFrame(
     np.random.randn(20, 3),
     columns=['a', 'b', 'c'])

st.line_chart(chart_data)

map_data = pd.DataFrame(
    np.random.randn(1000, 2) / [50, 50] + [37.76, -122.4],
    columns=['lat', 'lon'])

st.map(map_data)

x = st.slider('x')
st.write(x, 'squared is', x * x)

• 이런식으로 몇가지 기능만 갖고왔는데 데이터프레임을 보여준다던가 chart data로 표현한다던가 map 으로 table형식으로 보여주는등의 여러가지 기능이 있는데 document에 들어가면 더많은 기능들이 있다.
• 여기서 저는 Diffusion 의 Inpainting모델과 streamlit을 합쳐 각 모델을 올려두고 입력되는 사진에 따라서 결과물들을 보여주는 Demo 를 만들어보겠다.

 

최신 주요 기능

• 모델의 size가 클때 streamlit에서 실행시마다 모델을 reloading 하는 문제점이있는데 이 문제점을 해결하고 싶으면
  • st.cache_data
  • st.cache_resource 
  • 위 2가지 파이썬 데코레이터 기능을 사용해야한다.
• st.cache_data는 모델이 아닌 변하지 않는 데이터, 예를들어서 list의 값들이나 numpy array같은 곳에 모델에 전달할 변하지 않는 값들로써 사용할수있는데 이는 cache_data를 통해서 한번만 loading해두고 모델이 실행될때만 전달하게 해 memory leak과 cost를 줄일수 있다.
• st.cache_resource는 모델 파라미터를 저장하는데 사용하는데 여기서 loading한 모델은 무료 GPU VRAM 10G를 잡아먹는다. 이때도 cache_resource를 통해 한번만 loading하고 cost를 줄일수있다. flask로 구현해도 같은 방식을 적용할 수 있지만 server와 client를 구분해서 사용해야 한다는 단점이 있고  streamlit 프레임워크에 cache_resource를 사용하면 더 빠르게 사용가능하다. 

import streamlit as st
import PIL
import requests
import torch
from diffusers import AutoPipelineForInpainting, EulerAncestralDiscreteScheduler
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
from PIL import Image
import time
# Streamlit 앱의 타이틀 설정
st.title("얼굴부위별 시술설정")

@st.cache_resource
def load_model():
    try:
        model_id = "skin_treat_best"

        pipe = AutoPipelineForInpainting.from_pretrained(
            "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16 ,use_safetensors=True
        )
        # pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, safety_checker=None)
        # pipe.to("cpu")
        pipe.to('cuda') # GPU 있을때
        # pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        
        print('MODEL LOADING :OK')
    except Exception as e:
        print(e)
        print("MODEL LOAD : FAILE")
    return pipe
    
 pipe = load_model()

• 이런식으로 model을 load해 diffusion 의 pipe를 마련한다. 
• huggingface의 diffusion에서 pipe라는 방식을 채택한것은 pipe마다 구성요소들이 달라질수 있고 일정한 흐름 (DAG)방식으로 흐름이 이어져가기 때문에  pipe라는 이름을 붙여허 loading한다.
• 여기서 mask를 만드는것은 비밀이므로 mask를 생성하는 코드를 작성하거나 마스크 이미지 쌍을 갖고와서 실험해보길 바란다.

 

• Streamlit 에서 box option을 주고싶을때 select.box 를 사용해 지정시킬수 있다.

st.title("피부 치료 선택")

# 사용자에게 선택을 위한 드롭다운 메뉴 표시
selected_treatment = st.selectbox("치료 방법을 선택하세요:", list(instruction.keys()))

• 파일 업로드를 통해서 image를 받을수 있는 st객체를 만들고

uploaded_file = st.file_uploader("변환할 이미지 업로드", type=["jpg", "jpeg", "png"])

 

• 이런 형식으로 제출버튼을 만들고 클릭시에 model pipe에 넘기고 progress bar도 사용해서 변환되는 느낌처럼 만들수 도있다.

if st.button('변환'):
    if uploaded_file is not None and prompt:
        # PIL로 이미지 로드
        image = Image.open(uploaded_file).convert("RGB")
        h,w= image.size
        print(h,w)
        # image = image.resize((h//2,w//2))
        input_image = np.array(image)
        input = cv2.cvtColor(input_image,cv2.COLOR_BGR2RGB)

        h,w,c = input.shape
        mask_key = selected_treatment
        mask = draw_mask(input,mask_key)
        progress_bar = st.progress(0)
        status_text = st.empty()
        # 모델을 사용하여 이미지 변환
        # time.sleep(2)
        with torch.no_grad():
            outputs = pipe(prompt, image=image, mask_image=mask,num_inference_steps=30,strength=0.6,generator=generator,negative_prompt=negative_prompt)['images']

            for i in range(30):  # num_inference_steps에 해당하는 반복 횟수
                progress_bar.progress((i + 1) / 30)  # 프로그레스 바 업데이트
                status_text.text(f"변환 중... {int((i + 1) / 30 * 100)}%")

                time.sleep(0.01)  # 이 부분은 실제 모델 실행 시간에 맞게 조정해야 함
        print(outputs)
        st.image(image, caption="원본 이미지", use_column_width=False)
        # 변환된 이미지를 화면에 표시
        st.image(mask,caption='Mask 부위',use_column_width=False)
        st.image(outputs[0].resize((w,h)),caption='변환',use_column_width=False)

 

Result

• 1차적으로 이렇게 만들고

• 옵션 선택후

• 이런식으로 자기가 만든 모델의 output과 원하는 형태에 따라서 각기 Streamlit을 활용해서 Demo UI로 만들기 간편하고 Streamlit의 deploy기능을 사용하면 serving또한 간편하게 가능하다.