Generative AI Google Brain

open-models-serving-vertex-ai-pytorch-inference-paligemma-with-custom-handler

3 min read

Hugging Face DLCs: Serving PaliGemma using Pytorch Inference on Vertex AI with Custom Handler

Source notebook

Repo path: open-models/serving/vertex_ai_pytorch_inference_paligemma_with_custom_handler.ipynb · Open on GitHub · advanced

Deploys gated PaliGemma from Hugging Face to Vertex AI using a PyTorch DLC and custom handler.

Summary

The notebook teaches how to serve Google PaliGemma on Vertex AI with a Hugging Face PyTorch inference container and a custom handler. It walks through authentication, local handler testing, uploading model artifacts to Cloud Storage, registering the model, deploying to a GPU endpoint, and sending online predictions via Python, gcloud, and cURL.

Key code patterns

Initialize Vertex AI

import vertexai
 
vertexai.init(
    project=PROJECT_ID,
    location=LOCATION,
    staging_bucket=BUCKET_URI,
)

Configures the project, region, and staging bucket before model upload and deployment.

Load PaliGemma

processor = PaliGemmaProcessor.from_pretrained(
    "google/paligemma-3b-mix-448"
)
model = PaliGemmaForConditionalGeneration.from_pretrained(
    "google/paligemma-3b-mix-448",
    low_cpu_mem_usage=True,
    device_map="auto",
).eval()

Loads the gated Hugging Face model and processor for local inference testing.

Custom handler contract

class EndpointHandler:
    def __init__(self, model_dir="/opt/huggingface/model", **kwargs):
        self.processor = PaliGemmaProcessor.from_pretrained(model_dir)
        self.model = PaliGemmaForConditionalGeneration.from_pretrained(
            model_dir, device_map="auto", torch_dtype=torch.bfloat16
        ).eval()
 
    def __call__(self, data):
        return {"predictions": predictions}

Defines the Hugging Face inference handler used by the Vertex AI serving container.

Image prediction payload

prediction_request = {
    "instances": [{
        "prompt": "caption it",
        "image_base64": base64.b64encode(image_bytes).decode("utf-8"),
        "generation_kwargs": {"max_new_tokens": 100, "do_sample": False},
    }]
}

Shows the request shape expected by the custom handler for vision-language inference.

Upload and deploy model

model = Model.upload(
    display_name="google--paligemma-3b-mix-448",
    artifact_uri=str(model_uri),
    serving_container_image_uri=HF_DLC_IMAGE,
    serving_container_ports=[8080],
)
deployed_model = model.deploy(
    endpoint=Endpoint.create(display_name=ENDPOINT_NAME),
    machine_type="g2-standard-4",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
)

Registers model artifacts and deploys them to a GPU-backed Vertex AI endpoint.

Models & APIs used

  • Models: google/paligemma-3b-mix-448
  • APIs / services: Vertex AI, Cloud Storage, Artifact Registry
  • SDKs / libraries: vertexai, google-cloud-aiplatform, transformers, torch, huggingface_hub

When to use this

Use this pattern to serve a Hugging Face vision-language model on Vertex AI when standard container behavior needs custom preprocessing or postprocessing.

Gotchas & caveats

  • The PaliGemma model is gated and requires accepting the Hugging Face license plus a read-only access token.
  • The project must enable aiplatform.googleapis.com and artifactregistry.googleapis.com.
  • The service account needs Vertex AI User, Artifact Registry Reader, and Storage Object Admin roles.
  • LocalModel testing requires a local Docker installation.
  • Vertex AI deployment can take around 15 to 25 minutes.
  • The notebook uses us-central1 by default when GOOGLE_CLOUD_REGION is unset.
  • The deployed endpoint uses g2-standard-4 with one NVIDIA_L4 accelerator.

Best practices

  • Test the handler logic locally before packaging it as handler.py.
  • Use HF_HOME to control the local Hugging Face cache location.
  • Upload large model files to Cloud Storage with gsutil parallel composite uploads.
  • Validate request instances contain both prompt and image_base64.
  • Use torch.inference_mode() for generation.
  • Provide cleanup flags for endpoint, model, and bucket deletion.

Graph View

Backlinks