|
| 1 | +Deploying a PyTorch Stable Diffusion model as a Vertex AI Endpoint |
| 2 | +================================================================== |
| 3 | + |
| 4 | +Deploying large models, like Stable Diffusion, can be challenging and time-consuming. |
| 5 | + |
| 6 | +In this recipe, we will show how you can streamline the deployment of a PyTorch Stable Diffusion |
| 7 | +model by leveraging Vertex AI. |
| 8 | + |
| 9 | +PyTorch is the framework used by Stability AI on Stable |
| 10 | +Diffusion v1.5. Vertex AI is a fully-managed machine learning platform with tools and |
| 11 | +infrastructure designed to help ML practitioners accelerate and scale ML in production with |
| 12 | +the benefit of open-source frameworks like PyTorch. |
| 13 | + |
| 14 | +In four steps you can deploy a PyTorch Stable Diffusion model (v1.5). |
| 15 | + |
| 16 | +Deploying your Stable Diffusion model on a Vertex AI Endpoint can be done in four steps: |
| 17 | + |
| 18 | +* Create a custom TorchServe handler. |
| 19 | + |
| 20 | +* Upload model artifacts to Google Cloud Storage (GCS). |
| 21 | + |
| 22 | +* Create a Vertex AI model with the model artifacts and a prebuilt PyTorch container image. |
| 23 | + |
| 24 | +* Deploy the Vertex AI model onto an endpoint. |
| 25 | + |
| 26 | +Let’s have a look at each step in more detail. You can follow and implement the steps using the |
| 27 | +`Notebook example <https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/torchserve/dreambooth_stablediffusion.ipynb>`__. |
| 28 | + |
| 29 | +NOTE: Please keep in mind that this recipe requires a billable Vertex AI as explained in more details in the notebook example. |
| 30 | + |
| 31 | +Create a custom TorchServe handler |
| 32 | +---------------------------------- |
| 33 | + |
| 34 | +TorchServe is an easy and flexible tool for serving PyTorch models. The model deployed to Vertex AI |
| 35 | +uses TorchServe to handle requests and return responses from the model. |
| 36 | +You must create a custom TorchServe handler to include in the model artifacts uploaded to Vertex AI. Include the handler file in the |
| 37 | +directory with the other model artifacts, like this: `model_artifacts/handler.py`. |
| 38 | + |
| 39 | +After creating the handler file, you must package the handler as a model archiver (MAR) file. |
| 40 | +The output file must be named `model.mar`. |
| 41 | + |
| 42 | + |
| 43 | +.. code:: shell |
| 44 | +
|
| 45 | + !torch-model-archiver \ |
| 46 | + -f \ |
| 47 | + --model-name <your_model_name> \ |
| 48 | + --version 1.0 \ |
| 49 | + --handler model_artifacts/handler.py \ |
| 50 | + --export-path model_artifacts |
| 51 | +
|
| 52 | +Upload model artifacts to Google Cloud Storage (GCS) |
| 53 | +---------------------------------------------------- |
| 54 | + |
| 55 | +In this step we are uploading |
| 56 | +`model artifacts <https://github.com/pytorch/serve/tree/master/model-archiver#artifact-details>`__ |
| 57 | +to GCS, like the model file or handler. The advantage of storing your artifacts on GCS is that you can |
| 58 | +track the artifacts in a central bucket. |
| 59 | + |
| 60 | + |
| 61 | +.. code:: shell |
| 62 | +
|
| 63 | + BUCKET_NAME = "your-bucket-name-unique" # @param {type:"string"} |
| 64 | + BUCKET_URI = f"gs://{BUCKET_NAME}/" |
| 65 | +
|
| 66 | + # Will copy the artifacts into the bucket |
| 67 | + !gsutil cp -r model_artifacts $BUCKET_URI |
| 68 | +
|
| 69 | +Create a Vertex AI model with the model artifacts and a prebuilt PyTorch container image |
| 70 | +---------------------------------------------------------------------------------------- |
| 71 | + |
| 72 | +Once you've uploaded the model artifacts into a GCS bucket, you can upload your PyTorch model to |
| 73 | +`Vertex AI Model Registry <https://cloud.google.com/vertex-ai/docs/model-registry/introduction>`__. |
| 74 | +From the Vertex AI Model Registry, you have an overview of your models |
| 75 | +so you can better organize, track, and train new versions. For this you can use the |
| 76 | +`Vertex AI SDK <https://cloud.google.com/vertex-ai/docs/python-sdk/use-vertex-ai-python-sdk>`__ |
| 77 | +and this |
| 78 | +`pre-built PyTorch container <https://cloud.google.com/blog/products/ai-machine-learning/prebuilt-containers-with-pytorch-and-vertex-ai>`__. |
| 79 | + |
| 80 | + |
| 81 | +.. code:: shell |
| 82 | +
|
| 83 | + from google.cloud import aiplatform as vertexai |
| 84 | + PYTORCH_PREDICTION_IMAGE_URI = ( |
| 85 | + "us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-12:latest" |
| 86 | + ) |
| 87 | + MODEL_DISPLAY_NAME = "stable_diffusion_1_5-unique" |
| 88 | + MODEL_DESCRIPTION = "stable_diffusion_1_5 container" |
| 89 | +
|
| 90 | + vertexai.init(project='your_project', location='us-central1', staging_bucket=BUCKET_NAME) |
| 91 | +
|
| 92 | + model = aiplatform.Model.upload( |
| 93 | + display_name=MODEL_DISPLAY_NAME, |
| 94 | + description=MODEL_DESCRIPTION, |
| 95 | + serving_container_image_uri=PYTORCH_PREDICTION_IMAGE_URI, |
| 96 | + artifact_uri=BUCKET_URI, |
| 97 | + ) |
| 98 | +
|
| 99 | +Deploy the Vertex AI model onto an endpoint |
| 100 | +------------------------------------------- |
| 101 | + |
| 102 | +Once the model has been uploaded to Vertex AI Model Registry you can then take it and deploy |
| 103 | +it to an Vertex AI Endpoint. For this you can use the Console or the Vertex AI SDK. In this |
| 104 | +example you will deploy the model on a NVIDIA Tesla P100 GPU and n1-standard-8 machine. You can |
| 105 | +specify your machine type. |
| 106 | + |
| 107 | + |
| 108 | +.. code:: shell |
| 109 | +
|
| 110 | + endpoint = aiplatform.Endpoint.create(display_name=ENDPOINT_DISPLAY_NAME) |
| 111 | +
|
| 112 | + model.deploy( |
| 113 | + endpoint=endpoint, |
| 114 | + deployed_model_display_name=MODEL_DISPLAY_NAME, |
| 115 | + machine_type="n1-standard-8", |
| 116 | + accelerator_type="NVIDIA_TESLA_P100", |
| 117 | + accelerator_count=1, |
| 118 | + traffic_percentage=100, |
| 119 | + deploy_request_timeout=1200, |
| 120 | + sync=True, |
| 121 | + ) |
| 122 | +
|
| 123 | +If you follow this |
| 124 | +`notebook <https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/torchserve/dreambooth_stablediffusion.ipynb>`__ |
| 125 | +you can also get online predictions using the Vertex AI SDK as shown in the following snippet. |
| 126 | + |
| 127 | + |
| 128 | +.. code:: shell |
| 129 | +
|
| 130 | + instances = [{"prompt": "An examplePup dog with a baseball jersey."}] |
| 131 | + response = endpoint.predict(instances=instances) |
| 132 | +
|
| 133 | + with open("img.jpg", "wb") as g: |
| 134 | + g.write(base64.b64decode(response.predictions[0])) |
| 135 | +
|
| 136 | + display.Image("img.jpg") |
| 137 | +
|
| 138 | +Create a Vertex AI model with the model artifacts and a prebuilt PyTorch container image |
| 139 | + |
| 140 | +More resources |
| 141 | +-------------- |
| 142 | + |
| 143 | +This tutorial was created using the vendor documentation. To refer to the original documentation on the vendor site, please see |
| 144 | +`torchserve example <https://cloud.google.com/blog/products/ai-machine-learning/get-your-genai-model-going-in-four-easy-steps>`__. |
0 commit comments