Cloud Run GPU에서 TorchServe 및 Stable Diffusion을 실행하는 방법

1. 소개

개요

Cloud Run에 최근 GPU 지원이 추가되었습니다. 대기자 명단 공개 미리보기로 사용할 수 있습니다. 이 기능을 사용해 보고 싶다면 이 양식을 작성하여 대기자 명단에 등록하세요. Cloud Run은 클러스터를 관리할 필요 없이 컨테이너에서 코드를 간편하게 실행할 수 있는 Google Cloud의 컨테이너 플랫폼입니다.

현재 제공되는 GPU는 24GB의 vRAM이 있는 Nvidia L4 GPU입니다. Cloud Run 인스턴스당 GPU가 하나 있으며 Cloud Run 자동 확장은 계속 적용됩니다. 여기에는 최대 5개 인스턴스까지 수평 확장 (할당량 상향 조정 가능)과 요청이 없을 때 인스턴스를 0개로 축소하는 것이 포함됩니다.

이 Codelab에서는 안정적인 확산 XL을 사용하여 텍스트 프롬프트에서 이미지를 생성하는 TorchServe 앱을 만들고 배포합니다. 생성된 이미지는 base64로 인코딩된 문자열로 호출자에게 반환됩니다.

이 예는 Torchserve에서 Huggingface Diffusers를 사용하여 안정적인 확산 모델 실행을 기반으로 합니다. 이 Codelab에서는 Cloud Run에서 작동하도록 이 예시를 수정하는 방법을 보여줍니다.

학습할 내용

  • GPU를 사용하여 Cloud Run에서 Stable Diffusion XL 모델을 실행하는 방법

2. API 사용 설정 및 환경 변수 설정

이 Codelab을 사용하기 전에 몇 가지 API를 사용 설정해야 합니다. 이 Codelab에서는 다음 API를 사용해야 합니다. 다음 명령어를 실행하여 이러한 API를 사용 설정할 수 있습니다.

gcloud services enable run.googleapis.com \
    storage.googleapis.com \
    cloudbuild.googleapis.com \

그런 다음 이 Codelab 전체에서 사용할 환경 변수를 설정할 수 있습니다.

PROJECT_ID=<YOUR_PROJECT_ID>

REPOSITORY=repo
NETWORK_NAME=default
REGION=us-central1
IMAGE=us-central1-docker.pkg.dev/$PROJECT_ID/$REPOSITORY/gpu-torchserve

3. Torchserve 앱 만들기

먼저 소스 코드의 디렉터리를 만들고 해당 디렉터리로 이동합니다.

mkdir stable-diffusion-codelab && cd $_

config.properties 파일을 만듭니다. TorchServe의 구성 파일입니다.

inference_address=http://0.0.0.0:8080
enable_envvars_config=true
min_workers=1
max_workers=1
default_workers_per_model=1
default_response_timeout=1000
load_models=all
max_response_size=655350000
# to enable authorization, see https://github.com/pytorch/serve/blob/master/docs/token_authorization_api.md#how-to-set-and-disable-token-authorization
disable_token_authorization=true

이 예에서는 Cloud Run에서 작동하는 데 리슨 주소 http://0.0.0.0이 사용됩니다. Cloud Run의 기본 포트는 8080입니다.

requirements.txt 파일을 만듭니다.

python-dotenv
accelerate
transformers
diffusers
numpy
google-cloud-storage
nvgpu

stable_diffusion_handler.py라는 파일을 만듭니다.

from abc import ABC
import base64
import datetime
import io
import logging
import os

from diffusers import StableDiffusionXLImg2ImgPipeline
from diffusers import StableDiffusionXLPipeline
from google.cloud import storage
import numpy as np
from PIL import Image
import torch
from ts.torch_handler.base_handler import BaseHandler


logger = logging.getLogger(__name__)


def image_to_base64(image: Image.Image) -> str:
  """Convert a PIL image to a base64 string."""
  buffer = io.BytesIO()
  image.save(buffer, format="JPEG")
  image_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
  return image_str


class DiffusersHandler(BaseHandler, ABC):
  """Diffusers handler class for text to image generation."""

  def __init__(self):
    self.initialized = False

  def initialize(self, ctx):
    """In this initialize function, the Stable Diffusion model is loaded and

       initialized here.
    Args:
        ctx (context): It is a JSON Object containing information pertaining to
          the model artifacts parameters.
    """
    logger.info("Initialize DiffusersHandler")
    self.manifest = ctx.manifest
    properties = ctx.system_properties
    model_dir = properties.get("model_dir")
    model_name = os.environ["MODEL_NAME"]
    model_refiner = os.environ["MODEL_REFINER"]

    self.bucket = None

    logger.info(
        "GPU device count: %s",
        torch.cuda.device_count(),
    )
    logger.info(
        "select the GPU device, cuda is available: %s",
        torch.cuda.is_available(),
    )
    self.device = torch.device(
        "cuda:" + str(properties.get("gpu_id"))
        if torch.cuda.is_available() and properties.get("gpu_id") is not None
        else "cpu"
    )
    logger.info("Device used: %s", self.device)

    # open the pipeline to the inferenece model 
    # this is generating the image
    logger.info("Donwloading model %s", model_name)
    self.pipeline = StableDiffusionXLPipeline.from_pretrained(
        model_name,
        variant="fp16",
        torch_dtype=torch.float16,
        use_safetensors=True,
    ).to(self.device)
    logger.info("done donwloading model %s", model_name)

    # open the pipeline to the refiner
    # refiner is used to remove artifacts from the image
    logger.info("Donwloading refiner %s", model_refiner)
    self.refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        model_refiner,
        variant="fp16",
        torch_dtype=torch.float16,
        use_safetensors=True,
    ).to(self.device)
    logger.info("done donwloading refiner %s", model_refiner)

    self.n_steps = 40
    self.high_noise_frac = 0.8
    self.initialized = True
    # Commonly used basic negative prompts.
    logger.info("using negative_prompt")
    self.negative_prompt = ("worst quality, normal quality, low quality, low res, blurry")

  # this handles the user request
  def preprocess(self, requests):
    """Basic text preprocessing, of the user's prompt.

    Args:
        requests (str): The Input data in the form of text is passed on to the
          preprocess function.

    Returns:
        list : The preprocess function returns a list of prompts.
    """
    logger.info("Process request started")
    inputs = []
    for _, data in enumerate(requests):
      input_text = data.get("data")
      if input_text is None:
        input_text = data.get("body")
      if isinstance(input_text, (bytes, bytearray)):
        input_text = input_text.decode("utf-8")
      logger.info("Received text: '%s'", input_text)
      inputs.append(input_text)
    return inputs

  def inference(self, inputs):
    """Generates the image relevant to the received text.

    Args:
        input_batch (list): List of Text from the pre-process function is passed
          here

    Returns:
        list : It returns a list of the generate images for the input text
    """
    logger.info("Inference request started")
    # Handling inference for sequence_classification.
    image = self.pipeline(
        prompt=inputs,
        negative_prompt=self.negative_prompt,
        num_inference_steps=self.n_steps,
        denoising_end=self.high_noise_frac,
        output_type="latent",
    ).images
    logger.info("Done model")

    image = self.refiner(
        prompt=inputs,
        negative_prompt=self.negative_prompt,
        num_inference_steps=self.n_steps,
        denoising_start=self.high_noise_frac,
        image=image,
    ).images
    logger.info("Done refiner")

    return image

  def postprocess(self, inference_output):
    """Post Process Function converts the generated image into Torchserve readable format.

    Args:
        inference_output (list): It contains the generated image of the input
          text.

    Returns:
        (list): Returns a list of the images.
    """
    logger.info("Post process request started")
    images = []
    response_size = 0
    for image in inference_output:
      # Save image to GCS
      if self.bucket:
        image.save("temp.jpg")

        # Create a blob object
        blob = self.bucket.blob(
            datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + ".jpg"
        )

        # Upload the file
        blob.upload_from_filename("temp.jpg")

      # to see the image, encode to base64
      encoded = image_to_base64(image)
      response_size += len(encoded)
      images.append(encoded)

    logger.info("Images %d, response size: %d", len(images), response_size)
    return images

start.sh라는 파일을 만듭니다. 이 파일은 TorchServe를 시작하는 컨테이너의 진입점으로 사용됩니다.

#!/bin/bash

echo "starting the server"
# start the server. By default torchserve runs in backaround, and start.sh will immediately terminate when done
# so use --foreground to keep torchserve running in foreground while start.sh is running in a container  
torchserve --start --ts-config config.properties --models "stable_diffusion=${MAR_FILE_NAME}.mar" --model-store ${MAR_STORE_PATH} --foreground

그런 다음 다음 명령어를 실행하여 실행 파일로 만듭니다.

chmod 755 start.sh

dockerfile를 만듭니다.

# pick a version of torchserve to avoid any future breaking changes
# docker pull pytorch/torchserve:0.11.1-cpp-dev-gpu
FROM pytorch/torchserve:0.11.1-cpp-dev-gpu AS base

USER root

WORKDIR /home/model-server

COPY requirements.txt ./
RUN pip install --upgrade -r ./requirements.txt

# Stage 1 build the serving container.
FROM base AS serve-gcs

ENV MODEL_NAME='stabilityai/stable-diffusion-xl-base-1.0'
ENV MODEL_REFINER='stabilityai/stable-diffusion-xl-refiner-1.0'

ENV MAR_STORE_PATH='/home/model-server/model-store'
ENV MAR_FILE_NAME='model'
RUN mkdir -p $MAR_STORE_PATH

COPY config.properties ./
COPY stable_diffusion_handler.py ./
COPY start.sh ./

# creates the mar file used by torchserve
RUN torch-model-archiver --force --model-name ${MAR_FILE_NAME} --version 1.0 --handler stable_diffusion_handler.py -r requirements.txt --export-path ${MAR_STORE_PATH}

# entrypoint
CMD ["./start.sh"]

4. Cloud NAT 설정

Cloud NAT를 사용하면 인터넷에 액세스하고 HuggingFace에서 모델을 다운로드하는 데 더 높은 대역폭을 사용할 수 있으므로 배포 시간이 크게 단축됩니다.

Cloud NAT를 사용하려면 다음 명령어를 실행하여 Cloud NAT 인스턴스를 사용 설정합니다.

gcloud compute routers create nat-router --network $NETWORK_NAME --region us-central1
gcloud compute routers nats create vm-nat --router=nat-router --region=us-central1 --auto-allocate-nat-external-ips --nat-all-subnet-ip-ranges

5. Cloud Run 서비스 빌드 및 배포

Cloud Build에 코드를 제출합니다.

gcloud builds submit --tag $IMAGE

다음으로 Cloud Run에 배포합니다.

gcloud beta run deploy gpu-torchserve \
 --image=$IMAGE \
 --cpu=8 --memory=32Gi \
 --gpu=1 --no-cpu-throttling --gpu-type=nvidia-l4 \
 --allow-unauthenticated \
 --region us-central1 \
 --project $PROJECT_ID \
 --execution-environment=gen2 \
 --max-instances 1 \
 --network $NETWORK_NAME \
 --vpc-egress all-traffic

6. 서비스 테스트

다음 명령어를 실행하여 서비스를 테스트할 수 있습니다.

PROMPT_TEXT="a cat sitting in a magnolia tree"

SERVICE_URL=$(gcloud run services describe gpu-torchserve --region $REGION --format 'value(status.url)')

time curl $SERVICE_URL/predictions/stable_diffusion -d "data=$PROMPT_TEXT" | base64 --decode > image.jpg

현재 디렉터리에 image.jpg 파일이 표시됩니다. Cloud Shell 편집기에서 이미지를 열면 나무에 앉아 있는 고양이의 이미지를 볼 수 있습니다.

7. 축하합니다.

Codelab을 완료했습니다. 축하합니다.

Cloud Run GPU에 관한 문서를 검토하는 것이 좋습니다.

학습한 내용

  • GPU를 사용하여 Cloud Run에서 Stable Diffusion XL 모델을 실행하는 방법

8. 삭제

실수로 요금이 청구되지 않도록 하려면(예: 이 Cloud Run 작업이 무료 등급의 월별 Cloud Run 호출 할당량보다 실수로 더 많이 호출된 경우) Cloud Run 작업을 삭제하거나 2단계에서 만든 프로젝트를 삭제하면 됩니다.

Cloud Run 작업을 삭제하려면 https://console.cloud.google.com/run/에서 Cloud Run Cloud 콘솔로 이동하여 gpu-torchserve 서비스를 삭제합니다.

Cloud NAT 구성을 삭제하는 것도 좋습니다.

전체 프로젝트를 삭제하려면 https://console.cloud.google.com/cloud-resource-manager로 이동하여 2단계에서 만든 프로젝트를 선택하고 삭제를 선택합니다. 프로젝트를 삭제하면 Cloud SDK에서 프로젝트를 변경해야 합니다. gcloud projects list를 실행하여 사용 가능한 모든 프로젝트 목록을 볼 수 있습니다.