Colab 에서 streamlit app 실행하기

Colab 에서 streamlit app 실행하기

Colab 이란?

Google Colaboratory 서비스의 줄임말로, browser 를 통해 python을 작성/실행 가능합니다. Colab 공식 홈페이지 (https://colab.google/) 에서 다음과 같이 Colab 을 설명하고 있습니다.

별도의 설정이 필요 없는 호스팅형 Jupyter 노트북 서비스로, GPU와 TPU를 포함한 컴퓨팅 리소스에 무료로 액세스할 수 있습니다. Colab은 특히 머신 러닝, 데이터 과학 및 교육에 적합합니다.

위에서 언급한 것처럼 GPU, TPU 를 사용해볼 수 있어, 개인 PC 에서 머신 러닝을 실행해보기 어려운 경우 사용을 추천합니다.

 

Streamlit app 작성하기

Colab 에서 제공하는 TPU 를 이용해서 Generative AI 인 Stable Diffusion 을 실행하는 app 을 작성해보겠습니다.

 

Colab 환경 설정

"메뉴 >> 런타임 >> 런타임 유형 변경" 에서 아래와 같이 Python3, T4 GPU 를 선택하고,

런타임 유형 변경

 

아래 커맨드 실행해서 GPU 를 사용 가능한 것을 확인합니다.

!nvidia-smi

GPU 정보

 

아래 커맨드를 통해 개발에 필요한 package 를 설치해줍니다.

!npm install localtunnel
!pip install diffusers transformers scipy ftfy numpy
!pip install "ipywidgets>=7,<8"
!pip install streamlit

 

Hugging Face 에 로그인

Hugging Face 에서 Stable Diffusion 모델을 가지고 와서 실행하기 위해 Hugging Face 에 로그인을 해야 합니다.

from google.colab import output
from huggingface_hub import notebook_login

output.enable_custom_widget_manager()
notebook_login()

Hugging Face 로그인 화면

 

Hugging Face Access Token 을 입력하고 로그인을 합니다.

 

Streamlit app 작성하기

"runwayml/stable-diffusion-v1-5" 모델을 Hugging Face 로부터 받아오고, 사용자로부터 prompt 를 입력받아 그림을 생성하는 코드를 app.py 로 저장하는 것입니다. 

%%writefile app.py

import streamlit as st
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
from torch import autocast

MODEL = "runwayml/stable-diffusion-v1-5"

st.header("Streamlit Application")

pipe = StableDiffusionPipeline.from_pretrained(MODEL, revision="fp16", torch_dtype=torch.float16, use_auth_token=True)
pipe = pipe.to("cuda")

with st.form('form', clear_on_submit = True):
  prompt = st.text_input("Prompt: ", "")
  submitted = st.form_submit_button('Generate')

if submitted and prompt:
  prompt = [prompt]
  with autocast("cuda"):
    image = pipe(prompt, guidance_scale=8.0, num_inference_steps=25, height=512, width=512).images[0]

  st.image(image, use_column_width="auto", clamp=False, channels="RGB", output_format="auto")

위 코드를 실행하면 아래와 같이 prompt 를 입력받을 수 있는 화면이 보이는데, 다음 절에서 Colab 에서 실행한 Streamlit app 을 어떻게 확인할 수 있는지 알아보겠습니다.

실행 화면

 

Streamlit app 실행하기

!streamlit run /content/app.py &>/content/logs.txt &
!npx localtunnel --port 8501

저장된 app.py 를 실행하고, 외부에서 접근할 수 있도록 합니다.

 

localtunnel

Node.js 패키지로, 로컬 서버를 외부에서 접속할 수 있도록 해주는 도구입니다. 포트 포워딩과 같은 네트워크 설정 변경 없이 로컬 서버를 안전하게 외부에서 접근 가능하도록 해줍니다.

 

위 코드를 실행하면 localtunnel은 8051 포트를 통해 외부에서 접속할 수 있는 임시 주소를 생성합니다.

localtunnel

 

위에서 생성된 임시 주소(mean-kings-buy.loca.lt)를 클릭하면 아래와 같은 페이지가 열리는 것을 확인할 수 있습니다.

External URL 의 IP 주소 적는 창

 

Endpoint IP 는 /content/logs.txt 에 적혀있는 External URL 의 IP 주소(34.147.23.21)를 적어주면 됩니다. 

logs.txt

 

이미지 생성하기

Prompt 창에 원하는 글을 넣어 Generate 버튼을 클릭하면 이미지가 생성됩니다.

 

아래 내용을 Prompt 에 넣고 그림을 생성해보겠습니다.

"산은 설악산의 설악산이며, 해는 붉은색으로 지고 있습니다. 남자는 20대의 청년이며, 손에는 연필로 산의 풍경을 그리고 있습니다. "

Prompt 를 통해 생성된 이미지

 

지금까지 Stable Diffusion 을 실행할 수 있는 Streamlit app 을 작성하여 Colab에서 실행하는 방법에 대해 알아보았습니다.