notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4dc4391f6be7"
      },
      "source": [
        "# Vertex AI Model Garden - Hugging Face Local Inference\n",
        "\n",
        "<table><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_huggingface_local_inference.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4e8a0fdd6f44"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to install the necessary libraries and run local inference with various Hugging Face models in a [Colab Enterprise Instance](https://cloud.google.com/colab/docs).\n",
        "\n",
        "### Objective\n",
        "\n",
        "* Run local inference with various transformer or diffusion models.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69453bf7230e"
      },
      "source": [
        "## Install dependencies\n",
        "\n",
        "Before you begin, make sure you are connecting to a [Colab Enterprise runtime](https://cloud.google.com/colab/docs/connect-to-runtime) with GPU. If not, we recommend [creating a runtime template](https://cloud.google.com/colab/docs/create-runtime-template) with `g2-standard-16` machine type (or larger, see the descriptions of the model you want to try out below) to use `NVIDIA_L4` GPU. Then, [create a runtime](https://cloud.google.com/colab/docs/create-runtime) from that template.\n",
        "\n",
        "Some of the example models are [gated](https://huggingface.co/docs/hub/en/models-gated). If you want to try out one of such models, make sure to accept the model agreement from the Hugging Face model card page. Then, create a [Hugging Face read token](https://huggingface.co/docs/hub/en/security-tokens) and paste it below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3d342b32fb08"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "! pip3 install --upgrade pip\n",
        "! pip3 install xformers==0.0.25\n",
        "! pip3 install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1\n",
        "! pip3 install transformers~=4.44.2\n",
        "! pip3 install diffusers~=0.30.1\n",
        "! pip3 install accelerate~=0.33.0\n",
        "! pip3 install triton~=2.3.1\n",
        "! pip3 install tesseract~=0.1.3\n",
        "! pip3 install pytesseract~=0.3.13\n",
        "! apt-get update\n",
        "! apt-get install -y --no-install-recommends tesseract-ocr\n",
        "\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
        "os.environ[\"HF_HOME\"] = \"/content/hf-dir\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ad1a690839d5"
      },
      "source": [
        "## Sample code"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dgeuC_FrJ0ZY"
      },
      "source": [
        "#### [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) (Text-to-image)\n",
        "\n",
        "Generate images from a text description, which is also known as a `prompt`.\n",
        "\n",
        "This example runs [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) model with Diffusers [FluxPipeline](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux).\n",
        "\n",
        "GPU requirement:\n",
        "\n",
        "1. Running with **A100** GPU: Flux requires at least 30G memory in a single GPU. The inference takes ~4s. `a2-highgpu-1g` machine type is recommended.\n",
        "\n",
        "2. Running with **L4** GPU: If you don't have enough memory in a single GPU, you may select `enable_model_cpu_offload`. The pipeline then offloads the model to cpu with an overhead latency. The inference takes ~35s. `g2-standard-32` with 1 L4 GPU or `g2-standard-24` with 2 L4 GPU is recommended."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "pVkDyAHhJ0ZY"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "\n",
        "import torch\n",
        "from diffusers import FluxPipeline\n",
        "\n",
        "enable_model_cpu_offload = False  # @param {type:\"boolean\"}\n",
        "\n",
        "model_id = \"black-forest-labs/FLUX.1-schnell\"\n",
        "pipe = FluxPipeline.from_pretrained(\n",
        "    model_id,\n",
        "    torch_dtype=torch.bfloat16,\n",
        ")\n",
        "\n",
        "if not enable_model_cpu_offload:\n",
        "    pipe.to(\"cuda\")\n",
        "else:\n",
        "    pipe.enable_sequential_cpu_offload()\n",
        "    pipe.vae.enable_slicing()\n",
        "    pipe.vae.enable_tiling()\n",
        "    pipe.to(torch.float16)\n",
        "\n",
        "prompt = \"A cat holding a sign that says hello world\"  # @param {type:\"string\"}\n",
        "image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0]\n",
        "\n",
        "display(image)\n",
        "\n",
        "del pipe\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0a4008240483"
      },
      "source": [
        "#### [stabilityai/stable-diffusion-3-medium-diffusers](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) (Text-to-image)\n",
        "Generate photo-realistic images a text description, which is also known as a `prompt`.\n",
        "\n",
        "**This is a gated model. You need to agree the license displayed in the [model card](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers). Then, create a [Hugging Face read token](https://huggingface.co/docs/hub/en/security-tokens) and paste it in the HF_TOKEN field above.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5ec6b474d1be"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "\n",
        "import torch\n",
        "from diffusers import StableDiffusion3Pipeline\n",
        "\n",
        "model_id = \"stabilityai/stable-diffusion-3-medium-diffusers\"\n",
        "pipe = StableDiffusion3Pipeline.from_pretrained(\n",
        "    model_id,\n",
        "    torch_dtype=torch.float16,\n",
        "    token=HF_TOKEN,\n",
        ")\n",
        "pipe = pipe.to(\"cuda\")\n",
        "\n",
        "prompt = \"a photo of an astronaut riding a horse on mars\"\n",
        "image = pipe(prompt).images[0]\n",
        "\n",
        "display(image)\n",
        "\n",
        "del pipe\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ae94b9b23a52"
      },
      "source": [
        "#### [stabilityai/stable-diffusion-3-medium-diffusers](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) (Text guided image-to-image)\n",
        "Generate an image based on an initial image and a text prompt.\n",
        "\n",
        "**This is a gated model. You need to agree the license displayed in the [model card](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers). Then, create a [Hugging Face read token](https://huggingface.co/docs/hub/en/security-tokens) and paste it in the HF_TOKEN field above.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0acd70f4d08a"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "from io import BytesIO\n",
        "\n",
        "import requests\n",
        "import torch\n",
        "from diffusers import StableDiffusion3Img2ImgPipeline\n",
        "from PIL import Image\n",
        "\n",
        "device = \"cuda\"\n",
        "model_id_or_path = \"stabilityai/stable-diffusion-3-medium-diffusers\"\n",
        "pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(\n",
        "    model_id_or_path,\n",
        "    torch_dtype=torch.float16,\n",
        "    token=HF_TOKEN,\n",
        ")\n",
        "pipe = pipe.to(device)\n",
        "\n",
        "url = \"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg\"\n",
        "\n",
        "response = requests.get(url)\n",
        "init_image = Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
        "init_image = init_image.resize((768, 512))\n",
        "display(init_image)\n",
        "\n",
        "prompt = \"A fantasy landscape, trending on artstation\"\n",
        "\n",
        "images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images\n",
        "display(images[0])\n",
        "\n",
        "del pipe\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e76b3fe8d10c"
      },
      "source": [
        "#### [stabilityai/stable-diffusion-2-inpainting](https://huggingface.co/stabilityai/stable-diffusion-2-inpainting) (Image-inpainting)\n",
        "Generate an image based on an original image and prompt, only editing the areas denoted by a mask image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8bc3238be4e7"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "from io import BytesIO\n",
        "\n",
        "import requests\n",
        "import torch\n",
        "from diffusers import StableDiffusion3Img2ImgPipeline\n",
        "from PIL import Image\n",
        "\n",
        "device = \"cuda\"\n",
        "model_id_or_path = \"stabilityai/stable-diffusion-3-medium-diffusers\"\n",
        "pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(\n",
        "    model_id_or_path,\n",
        "    torch_dtype=torch.float16,\n",
        "    token=HF_TOKEN,\n",
        ")\n",
        "pipe = pipe.to(device)\n",
        "\n",
        "url = \"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg\"\n",
        "\n",
        "response = requests.get(url)\n",
        "init_image = Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
        "init_image = init_image.resize((768, 512))\n",
        "display(init_image)\n",
        "\n",
        "prompt = \"A fantasy landscape, trending on artstation\"\n",
        "\n",
        "images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images\n",
        "display(images[0])\n",
        "\n",
        "del pipe\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ade95a9b20e"
      },
      "source": [
        "#### [impira/layoutlm-document-qa](https://huggingface.co/impira/layoutlm-document-qa) (Document question answering)\n",
        "Answer questions about a given document."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "492d9f1de3f2"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "\n",
        "from transformers import pipeline\n",
        "\n",
        "nlp = pipeline(\n",
        "    \"document-question-answering\",\n",
        "    model=\"impira/layoutlm-document-qa\",\n",
        ")\n",
        "\n",
        "print(\n",
        "    nlp(\n",
        "        \"https://templates.invoicehome.com/invoice-template-us-neat-750px.png\",\n",
        "        \"What is the invoice number?\",\n",
        "    )\n",
        ")\n",
        "# [{'score': 0.4251753091812134, 'answer': 'us-001', 'start': 16, 'end': 16}]\n",
        "\n",
        "print(\n",
        "    nlp(\n",
        "        \"https://miro.medium.com/max/787/1*iECQRIiOGTmEFLdWkVIH2g.jpeg\",\n",
        "        \"What is the purchase amount?\",\n",
        "    )\n",
        ")\n",
        "# [{'score': 0.999853253364563, 'answer': '$1,000,000,000', 'start': 97, 'end': 97}]\n",
        "\n",
        "print(\n",
        "    nlp(\n",
        "        \"https://www.accountingcoach.com/wp-content/uploads/2013/10/income-statement-example@2x.png\",\n",
        "        \"What are the 2020 net sales?\",\n",
        "    )\n",
        ")\n",
        "# [{'score': 0.9726569652557373, 'answer': '$ 3,980', 'start': 11, 'end': 12}]\n",
        "\n",
        "del nlp\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lxXSN_HAX7QL"
      },
      "source": [
        "#### [Alibaba-NLP/gte-large-en-v1.5](https://huggingface.co/Alibaba-NLP/gte-large-en-v1.5) (Feature-extraction)\n",
        "Get text embeddings from a sentence."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Q29z6kOdX7QL"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "\n",
        "from transformers import AutoTokenizer, pipeline\n",
        "\n",
        "model_path = \"Alibaba-NLP/gte-large-en-v1.5\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
        "feature_extraction = pipeline(\n",
        "    \"feature-extraction\",\n",
        "    model=model_path,\n",
        "    tokenizer=tokenizer,\n",
        "    trust_remote_code=True,\n",
        ")\n",
        "features = feature_extraction(\"i am sentence\")\n",
        "\n",
        "print(features[0])\n",
        "\n",
        "del feature_extraction, tokenizer\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pdTm_rNeX7QL"
      },
      "source": [
        "#### [google/gemma-2-2b](https://huggingface.co/google/gemma-2-2b) (Text-generation)\n",
        "Generate text from another text; For example, fill in incomplete text or paraphrase.\n",
        "\n",
        "**This is a gated model. You need to agree the license displayed in the [model card](https://huggingface.co/google/gemma-2-2b). Then, create a [Hugging Face read token](https://huggingface.co/docs/hub/en/security-tokens) and paste it in the HF_TOKEN field above.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KkLEv7VlX7QL"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "\n",
        "from transformers import pipeline\n",
        "\n",
        "pipe = pipeline(\n",
        "    \"text-generation\",\n",
        "    model=\"google/gemma-2-2b\",\n",
        "    device=\"cuda\",\n",
        "    token=HF_TOKEN,\n",
        ")\n",
        "\n",
        "text = \"Once upon a time,\"\n",
        "outputs = pipe(text, max_new_tokens=256)\n",
        "response = outputs[0][\"generated_text\"]\n",
        "print(response)\n",
        "\n",
        "del pipe\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_huggingface_local_inference.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}