About this codelab
1. Introduction
Overview
Cloud Run recently added GPU support. It's available as a waitlisted public preview. If you're interested in trying out the feature, fill out this form to join the waitlist. Cloud Run is a container platform on Google Cloud that makes it straightforward to run your code in a container, without requiring you to manage a cluster.
Today, the GPUs we make available are Nvidia L4 GPUs with 24 GB of vRAM. There's one GPU per Cloud Run instance, and Cloud Run auto scaling still applies. That includes scaling out up to 5 instances (with quota increase available), as well as scaling down to zero instances when there are no requests.
In this codelab, you will create and deploy a TorchServe app that uses stable diffusion XL to generate images from a text prompt. The generated image is returned to the caller as a base64 encoded string.
This example is based on Running Stable diffusion model using Huggingface Diffusers in Torchserve. This codelab shows you how to modify this example to work with Cloud Run.
What you'll learn
- How to run a Stable Diffusion XL model on Cloud Run using GPUs
2. Enable APIs and Set Environment Variables
Before you can start using this codelab, there are several APIs you will need to enable. This codelab requires using the following APIs. You can enable those APIs by running the following command:
gcloud services enable run.googleapis.com \ storage.googleapis.com \ cloudbuild.googleapis.com \
Then you can set environment variables that will be used throughout this 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. Create the Torchserve app
First, create a directory for the source code and cd into that directory.
mkdir stable-diffusion-codelab && cd $_
Create a config.properties
file. This is the configuration file for 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
Note that in this example, the listening address http://0.0.0.0 is used to work on Cloud Run. The default port for Cloud Run is port 8080.
Create a requirements.txt
file.
python-dotenv accelerate transformers diffusers numpy google-cloud-storage nvgpu
Create a file called 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
Create a file called start.sh
This file is used as an entrypoint in the container to start 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
And then run the following command to make it an executable file.
chmod 755 start.sh
Create a 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. Setup Cloud NAT
Cloud NAT allows you to have higher bandwidth to access the internet and download the model from HuggingFace, which will significantly speed up your deployment times.
To use Cloud NAT, run the following command to enable a Cloud NAT instance:
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. Build and deploy the Cloud Run service
Submit your code to Cloud Build.
gcloud builds submit --tag $IMAGE
Next, deploy to 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. Test the service
You can test the service by running the following commands:
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
You will see image.jpg
file appear in your current directory. You can open the image in the Cloud Shell Editor to see the image of a cat sitting in a tree.
7. Congratulations!
Congratulations for completing the codelab!
We recommend reviewing the documentation on Cloud Run GPUs.
What we've covered
- How to run a Stable Diffusion XL model on Cloud Run using GPUs
8. Clean up
To avoid inadvertent charges, (for example, if this Cloud Run job is inadvertently invoked more times than your monthly Cloud Run invokement allocation in the free tier), you can either delete the Cloud Run job or delete the project you created in Step 2.
To delete the Cloud Run job, go to the Cloud Run Cloud Console at https://console.cloud.google.com/run/ and delete the gpu-torchserve
service.
You will also want to delete your Cloud NAT configuration.
If you choose to delete the entire project, you can go to https://console.cloud.google.com/cloud-resource-manager, select the project you created in Step 2, and choose Delete. If you delete the project, you'll need to change projects in your Cloud SDK. You can view the list of all available projects by running gcloud projects list
.