debug Meghan reported failure

* feat: SDK2 remote training

* fix: 3.10

* fix: 3.10

* fix: py check

* fix: py check

* fix: review comments

community-content/CODEOWNERS

@@ -12,4 +12,12 @@
 /prediction_featurestore_integration @googleapis/vertex-prediction-team
 /vertex_vision_model_garden/model_oss/util @weigary
 /vertex_vision_model_garden/model_oss/diffusers @weigary
+/vertex_vision_model_garden/model_oss/keras @dstnluong-google
 /vertex_vision_model_garden/model_oss/transformers @dstnluong-google
+/vertex_vision_model_garden/model_oss/pic2word @jismailyan-google
+/vertex_vision_model_garden/model_oss/open_clip @lydhr
+/vertex_vision_model_garden/model_oss/movinet @KCFindstr
+/vertex_vision_model_garden/model_oss/data_converter @KCFindstr
+/vertex_vision_model_garden/model_oss/peft @weigary
+/vertex_vision_model_garden/model_oss/lm-evaluation-harness @kathyyu-google
+

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -9,7 +9,7 @@ binarize_column_using_Pandas_on_CSV_data_op = components.load_component_from_url
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -23,7 +23,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 #train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -8,7 +8,7 @@ fill_all_missing_values_using_Pandas_on_CSV_data_op = components.load_component_
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -22,7 +22,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml

@@ -64,8 +64,8 @@ implementation:
           labels["component-source"] = "github-com-ark-kun-pipeline-components"
 
           # The serving container decides the model type based on the model file extension.
-          # So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.pkl
-          _, renamed_model_path = tempfile.mkstemp(suffix=".pkl")
+          # So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.bst
+          _, renamed_model_path = tempfile.mkstemp(suffix=".bst")
           shutil.copyfile(src=model_path, dst=renamed_model_path)
 
           model = aiplatform.Model.upload_xgboost_model_file(

community-content/pipeline_components/image_ml_model_training/load_image_classification_model/component.yaml

@@ -87,7 +87,7 @@ outputs:
 - {name: image_size_path, type: HeightWidth}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.
@@ -109,4 +109,4 @@ implementation:
       {inputValue: l2_regularization_penalty},
       --image-size-path,
       {outputPath: image_size_path},
-    ]
+    ]

community-content/pipeline_components/image_ml_model_training/preprocess_image_data/component.yaml

@@ -34,7 +34,7 @@ outputs:
     path for the validation data,'}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/train_image_classification_model/component.yaml

@@ -55,7 +55,7 @@ outputs:
     for the saved model,'}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_csv/component.yaml

@@ -20,7 +20,7 @@ outputs:
     path for the TFRecord image data}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_jsonl/component.yaml

@@ -22,7 +22,7 @@ outputs:
     path for the TFRecord image data}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/vertex_vision_model_garden/model_oss/data_converter/common_lib.py

@@ -0,0 +1,623 @@
+"""Library with functions to use for data conversion."""
+
+import json
+import os
+import random
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+import uuid
+
+from absl import logging
+import apache_beam as beam
+import cv2
+import numpy as np
+import pandas as pd
+import PIL
+from PIL import Image
+import tensorflow as tf
+import yaml
+
+from util import constants
+from util import fileutils
+from apache_beam.options import pipeline_options
+
+REFORMATTED_CSV_SUFFIX = '-reformatted.csv'
+
+LABEL_MAP_NAME = 'label_map.yaml'
+
+_SPLIT_RATIO_ERROR_THRESHOLD = 1e-5
+# Internal constant. Only for distinguishing rows without ML use.
+ML_USE_UNASSIGNED = 'unassigned'
+ALL_ML_USES = (
+    constants.ML_USE_TRAINING,
+    constants.ML_USE_VALIDATION,
+    constants.ML_USE_TEST,
+    ML_USE_UNASSIGNED,
+)
+COLUMN_NAME_ML_USE = 'ml_use'
+COLUMN_NAME_GCS_FILE_PATH = 'gcs_file_path'
+COLUMN_NAME_LABEL = 'label'
+COLUMN_NAME_START_SEC = 'start_sec'
+COLUMN_NAME_END_SEC = 'end_sec'
+# Output filenames
+TRAIN_TFRECORD_NAME = 'train.tfrecord'
+VALIDATION_TFRECORD_NAME = 'val.tfrecord'
+TEST_TFRECORD_NAME = 'test.tfrecord'
+# Jsonl keys
+JSON_GCS_URI_KEY = 'imageGcsUri'
+JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels'
+JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use'
+# I/O parameters
+READ_CHUNK_SIZE = 1024 * 1024 * 1024  # 1GB
+
+
+class WriteToTFRecord(beam.DoFn):
+  """DoFn to write TF examples to sharded TF record files."""
+
+  def __init__(
+      self,
+      output_prefix: str,
+      num_shards: int,
+      convert_fn: Callable[[Dict[str, Any]], tf.train.Example],
+  ):
+    self.output_prefix = output_prefix
+    self.num_shards = num_shards
+    self.writer: list[tf.io.TFRecordWriter] = []
+    self.sharded_files: list[str] = []
+    self.convert_fn = convert_fn
+    self.success_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Success'
+    )
+    self.failure_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Failure'
+    )
+
+  def start_bundle(self):
+    logging.info('Start writing TF Record to %s.', self.output_prefix)
+    unique_str = uuid.uuid4().hex
+    for i in range(self.num_shards):
+      uri = f'{self.output_prefix}-{i}-{unique_str}'
+      self.sharded_files.append(uri)
+      self.writer.append(tf.io.TFRecordWriter(uri))
+
+  def process(self, data: Dict[str, Any]) -> Iterable[Tuple[int, str]]:
+    try:
+      example = self.convert_fn(data)
+      data = example.SerializeToString()
+      idx = hash(data) % self.num_shards
+      self.writer[idx].write(data)
+      self.success_counter.inc()
+      yield (idx, self.sharded_files[idx])
+    # pylint: disable-next=broad-exception-caught
+    except Exception as err:
+      logging.error('Failed to process %s', data)
+      logging.exception(err)
+      self.failure_counter.inc()
+
+  def finish_bundle(self):
+    logging.info('Finish writing TF Record to %s.', self.output_prefix)
+    for writer in self.writer:
+      writer.close()
+    self.writer = []
+
+
+def convert_to_feature(
+    value: Union[List[Union[int, float, bytes]], int, float, bytes],
+    value_type: Optional[str] = None,
+) -> tf.train.Feature:
+  """Converts the given python object to a tf.train.Feature.
+
+  This is copied from tensorflow_models/official/vision/data/tfrecord_lib.py.
+
+  Args:
+    value: int, float, bytes or a list of them.
+    value_type: optional, if specified, forces the feature to be of the given
+      type. Otherwise, type is inferred automatically. Can be one of ['bytes',
+      'int64', 'float', 'bytes_list', 'int64_list', 'float_list']
+
+  Returns:
+    feature: A tf.train.Feature object.
+  """
+
+  if value_type is None:
+    element = value[0] if isinstance(value, list) else value
+
+    if isinstance(element, bytes):
+      value_type = 'bytes'
+
+    elif isinstance(element, (int, np.integer)):
+      value_type = 'int64'
+
+    elif isinstance(element, (float, np.floating)):
+      value_type = 'float'
+
+    else:
+      raise ValueError(
+          'Cannot convert type {} to feature'.format(type(element))
+      )
+
+    if isinstance(value, list):
+      value_type = value_type + '_list'
+
+  if value_type == 'int64':
+    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
+
+  elif value_type == 'int64_list':
+    value = np.asarray(value).astype(np.int64).reshape(-1)
+    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
+
+  elif value_type == 'float':
+    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
+
+  elif value_type == 'float_list':
+    value = np.asarray(value).astype(np.float32).reshape(-1)
+    return tf.train.Feature(float_list=tf.train.FloatList(value=value))
+
+  elif value_type == 'bytes':
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
+
+  elif value_type == 'bytes_list':
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+
+  else:
+    raise ValueError('Unknown value_type parameter - {}'.format(value_type))
+
+
+def convert_to_string_feature(
+    value: str, encoding: str = 'utf-8'
+) -> tf.train.Feature:
+  """Returns a bytes_list from an encoded string."""
+  return convert_to_feature(value.encode(encoding))
+
+
+def convert_to_list_string_feature(
+    lst: list[str], encoding: str = 'utf-8'
+) -> tf.train.Feature:
+  """Returns a bytes_list from a list of encoded strings."""
+  return convert_to_feature([value.encode(encoding) for value in lst])
+
+
+def create_ml_use_array_with_split(
+    total_size: int,
+    split_ratio: Sequence[float],
+) -> list[str]:
+  """Create randomized list of 'training', 'validation', 'test'.
+
+  The list of will be of length total_size with ratios according to train_size,
+  validation_size, and test_size.
+
+  Args:
+    total_size: Length of sequence to return
+    split_ratio: Proportions to split into 'training', 'validation', and 'test'
+
+  Returns:
+    List containing 'training', 'validation', and 'test'
+  """
+  train_size, validation_size, _ = split_ratio
+  num_train = round(train_size * total_size)
+  num_validation = round(validation_size * total_size)
+  num_test = total_size - num_train - num_validation
+  ml_use_row = (
+      [constants.ML_USE_TRAINING] * num_train
+      + [constants.ML_USE_VALIDATION] * num_validation
+      + [constants.ML_USE_TEST] * num_test
+  )
+  random.shuffle(ml_use_row)
+  return ml_use_row
+
+
+def format_ml_use_column(df: pd.DataFrame):
+  df[COLUMN_NAME_ML_USE].replace(
+      # We need to support non-standard ML uses other than documented ones,
+      # since they are used by some existing datasets.
+      [r'(?i)^train(ing)?$', r'(?i)^test$', r'(?i)^validat(ion|e)$'],
+      [
+          constants.ML_USE_TRAINING,
+          constants.ML_USE_TEST,
+          constants.ML_USE_VALIDATION,
+      ],
+      inplace=True,
+      regex=True,
+  )
+
+
+def insert_missing_ml_use(df: pd.DataFrame) -> None:
+  """For every row that does not have ml_use as the first column, insert a column containing 'unassigned' to the front.
+
+  Args:
+    df: The DataFrame to process. The first column should be 'ml_use'.
+  """
+  df[COLUMN_NAME_ML_USE].fillna(ML_USE_UNASSIGNED, inplace=True)
+  rows_to_fill = ~df[COLUMN_NAME_ML_USE].isin(ALL_ML_USES)
+  df.loc[rows_to_fill] = df[rows_to_fill].shift(
+      axis=1, fill_value=ML_USE_UNASSIGNED
+  )
+
+
+def replace_unassigned_ml_use(
+    ml_uses: List[str],
+    split_ratio: Sequence[float],
+):
+  """Replace `unassigned` in ml_uses with `training`, `validation`, and `test` with ratios according to split_ratio.
+
+  Args:
+    ml_uses: List of ml_use string values.
+    split_ratio: Proportions to split into `training`, `validation`, and `test`.
+  """
+  unassigned_indices = [
+      i for i, ml_use in enumerate(ml_uses) if ml_use == ML_USE_UNASSIGNED
+  ]
+  ml_use_arr = create_ml_use_array_with_split(
+      len(unassigned_indices), split_ratio
+  )
+  for unassigned_index, ml_use in zip(unassigned_indices, ml_use_arr):
+    ml_uses[unassigned_index] = ml_use
+
+
+def merge_seq_into_dicts(
+    key: str, values: Sequence[Any], dicts: Sequence[Dict[Any, Any]]
+):
+  """Merges a list of values into a list of dicts, inserted with the given key.
+
+  Args:
+    key: Key to insert or overwrite in the dictionary.
+    values: A list of values to insert.
+    dicts: A list of dictionaries. Each value will be inserted into the
+      corresponding dictionary. The original value will be overwritten if the
+      key already existed.
+
+  Raises:
+    ValueError: The values and dicts have different lengths.
+  """
+  if len(values) != len(dicts):
+    raise ValueError(
+        f'Length of values and dicts must match, got {len(values)} and'
+        f' {len(dicts)}'
+    )
+  for val, d in zip(values, dicts):
+    d[key] = val
+
+
+def drop_invalid_rows(df: pd.DataFrame) -> int:
+  """Drops DataFrame rows missing the gcs_file_path column or the label column.
+
+  Args:
+    df: The DataFrame to process in place.
+
+  Returns:
+    The number of rows dropped.
+  """
+  original_rows = df.shape[0]
+  df.dropna(subset=[COLUMN_NAME_GCS_FILE_PATH, COLUMN_NAME_LABEL], inplace=True)
+  dropped_num = original_rows - df.shape[0]
+  if dropped_num > 0:
+    df.reset_index(drop=True, inplace=True)
+  return dropped_num
+
+
+def check_split_ratio(split_ratio: Sequence[float]):
+  """Checks if the give split ratio is valid.
+
+  Args:
+    split_ratio: Proportions to split into 'training', 'validation', and 'test'
+
+  Raises:
+    ValueError: Must have valid entries, correct length, and sum to 1.
+  """
+  if len(split_ratio) != 3:
+    raise ValueError('split_ratio must contain exactly 3 values.')
+  if abs(sum(split_ratio) - 1) > _SPLIT_RATIO_ERROR_THRESHOLD:
+    raise ValueError('split_ratio must sum to 1.')
+  if not all([0 <= val <= 1 for val in split_ratio]):
+    raise ValueError('Entries of split_ratio must be in the range [0, 1].')
+
+
+def check_num_shard(num_shard: Sequence[int]):
+  """Checks if the number of shards is valid.
+
+  Args:
+    num_shard: The number of shards for each tfrecord.
+
+  Raises:
+    ValueError: Must have valid entries and correct length.
+  """
+  if len(num_shard) != 3:
+    raise ValueError('num_shard must contain exactly 3 values.')
+  if not all([val >= 1 for val in num_shard]):
+    raise ValueError('Shards must be at least 1.')
+
+
+def create_label_map_yaml(meta_data_path: str, output_dir: str) -> None:
+  """Generate label_map.yaml from meta_data.yaml.
+
+  Args:
+    meta_data_path: Path to a meta_data.yaml file.
+    output_dir: Directory to output label_map.yaml.
+  """
+  tf.io.gfile.copy(
+      meta_data_path, os.path.join(output_dir, LABEL_MAP_NAME), overwrite=True
+  )
+
+
+def reformat_bbox(
+    bbox: Sequence[int], img_width: int, img_height: int
+) -> Tuple[float, float, float, float]:
+  """Converts XYWH unnormalized bounding box with to a normalized XYXY bounding box.
+
+  Args:
+    bbox: Relative bounding box with unnormalized coordinates as [x, y, width,
+      height].
+    img_width: Image's pixel width.
+    img_height: Image's pixel height.
+
+  Returns:
+    Absolute bounding box with normalized coordinates as
+      [xmin, ymin, xmax, ymax].
+  """
+  x, y, width, height = bbox
+  xmin = x / img_width
+  ymin = y / img_height
+  xmax = (x + width) / img_width
+  ymax = (y + height) / img_height
+  return xmin, ymin, xmax, ymax
+
+
+def encode_image(
+    filepath: str,
+    output_shape: Optional[Sequence[int]] = None,
+    image_format: str = 'png',
+) -> Tuple[bytes, Sequence[int]]:
+  """Encodes an image at the given path.
+
+  Args:
+    filepath: Path to the image.
+    output_shape: The output shape of the image, (height, width).
+    image_format: The format of the output image.
+
+  Returns:
+    The encoded image data in bytes and the shape of the image, (height, width).
+
+  Raises:
+    IOError: The image file is corrupt.
+  """
+  filepath = fileutils.force_gcs_fuse_path(filepath)
+  with open(filepath, 'rb') as f:
+    # If an output_shape is specified, resize the image and set data to the new
+    # bytes.
+    try:
+      img = Image.open(f)
+    except PIL.UnidentifiedImageError as e:
+      raise IOError(f'Failed to open {filepath}') from e
+
+    try:
+      if output_shape is not None:
+        rgb_img = img.resize((output_shape[1], output_shape[0])).convert('RGB')
+      else:
+        rgb_img = img.convert('RGB')
+      rgb_img = np.array(rgb_img)
+
+      _, data = cv2.imencode(f'.{image_format}', rgb_img)
+      data = data.tobytes()
+      return data, rgb_img.shape
+    except cv2.error as e:
+      raise IOError(f'Failed to encode {filepath}') from e
+    finally:
+      img.close()
+
+
+def encode_video(
+    filepath: str,
+    start_sec: float,
+    end_sec: float,
+    output_fps: int = 5,
+    output_shape: Optional[Sequence[int]] = None,
+    image_format: str = 'jpg',
+) -> Sequence[bytes]:
+  """Encodes a video clip at the given path with start and end timestamps.
+
+  Args:
+    filepath: Path to the video.
+    start_sec: Start timestamp of the video clip in seconds.
+    end_sec: End timestamp of the video clip in seconds.
+    output_fps: The output frame rate per second.
+    output_shape: The output shape of each frame, (height, width).
+    image_format: The format of the encoded frames.
+
+  Returns:
+    A list of the encoded frames data in bytes.
+
+  Raises:
+    IOError if the video file is corrupt.
+  """
+  filepath = fileutils.force_gcs_fuse_path(filepath)
+  video = None
+
+  try:
+    video = cv2.VideoCapture(filepath)
+    frames = []
+    frame_interval = 1 / output_fps
+    total_frames = video.get(cv2.CAP_PROP_FRAME_COUNT)
+    original_fps = video.get(cv2.CAP_PROP_FPS)
+    if not original_fps:
+      # 0 or None indicates the video is invalid
+      raise IOError(f'Failed to load {filepath}')
+    video_length = total_frames / original_fps
+    start_sec = max(start_sec, 0)
+    end_sec = min(end_sec, video_length)
+    for t in np.arange(start_sec, end_sec, frame_interval):
+      frame_idx = min(total_frames - 1, round(t * original_fps))
+      video.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+      ret, frame = video.read()
+      if not ret:
+        raise IOError(f'Failed to load {filepath} at frame {frame_idx}')
+      if output_shape is not None:
+        frame = cv2.resize(frame, (output_shape[1], output_shape[0]))
+      _, data = cv2.imencode(f'.{image_format}', frame)
+      frames.append(data.tobytes())
+  except cv2.error as e:
+    raise IOError(f'Failed to load {filepath}') from e
+  finally:
+    if video:
+      video.release()
+  return frames
+
+
+def create_label_map(
+    labels: Sequence[str],
+) -> Tuple[Sequence[int], Dict[int, str]]:
+  """Creates a label map from a sequence of label strings.
+
+  Args:
+    labels: The sequence of labels to create label map from. Must not contain
+      invalid values, which means data without labels should be filtered first.
+
+  Returns:
+    The integer labels and the mapping from integers to the original strings.
+  """
+  inverse_label_map: Dict[str, int] = dict()
+  num_labels = 0
+  for label in labels:
+    if label not in inverse_label_map:
+      num_labels += 1
+      inverse_label_map[label] = num_labels
+  int_labels = [inverse_label_map[label] for label in labels]
+  label_map = {value: key for key, value in inverse_label_map.items()}
+  return int_labels, label_map
+
+
+def write_label_map(output_file: str, label_map: Dict[int, str]) -> None:
+  """Writes a label map to the output file, which can be a GCS uri."""
+  with tf.io.gfile.GFile(output_file, 'w') as f:
+    yaml.dump({'label_map': label_map}, f)
+
+
+def detectron_json_to_image_rows(input_json: str) -> list[Dict[str, Any]]:
+  """Converts a Detectron JSON file to a list of image rows.
+
+  Args:
+    input_json: A path to a Detectron JSON or JSONL file.
+
+  Returns:
+    A list of dictionaries, where each dictionary contains Detectron format
+    entry.
+
+  Raises:
+    ValueError: If the input JSON is invalid.
+  """
+
+  image_rows = []
+  with tf.io.gfile.GFile(input_json, 'r') as f:
+    for line in f:
+      json_data = json.loads(line)
+      if isinstance(json_data, dict):
+        image_rows.append(json_data)
+      elif isinstance(json_data, list):
+        image_rows.extend(json_data)
+      else:
+        raise ValueError(
+            'The input JSON is invalid. Dict or list is expected, but got '
+            f'{type(json_data)}.'
+        )
+  return image_rows
+
+
+def coco_json_to_image_rows(
+    input_json: str,
+) -> List[Dict[str, Any]]:
+  """Converts a COCO JSON file to a list of image rows.
+
+  Args:
+    input_json: A path to a COCO JSON or JSONL file.
+
+  Returns:
+    A list of dictionaries, where each dictionary contains COCO format entry.
+
+  Raises:
+    ValueError: If the input JSON is invalid.
+  """
+
+  with tf.io.gfile.GFile(input_json, 'r') as f:
+    coco_json = json.load(f)
+  if 'annotations' not in coco_json:
+    raise ValueError('"annotations" is not in the dataset.')
+  if 'images' not in coco_json:
+    raise ValueError('"images" is not in the dataset.')
+
+  images = coco_json['images']
+  return images
+
+
+def partition_by_ml_use(element: Dict[str, Any], num_partitions: int) -> int:
+  """Beam partition function to split data by ml_use."""
+  del num_partitions
+  try:
+    partition = ALL_ML_USES.index(element[COLUMN_NAME_ML_USE])
+  except Exception as e:
+    raise ValueError(f'Invalid ML use: {element[COLUMN_NAME_ML_USE]}') from e
+  return partition
+
+
+def run_beam_pipeline(pipeline: Any) -> None:
+  """Runs a beam pipeline. Works in both internal and docker environment."""
+  options = pipeline_options.PipelineOptions([
+      '--runner=FlinkRunner',
+      '--faster_copy',
+      '--max_parallelism', '8',
+  ])
+  p = beam.Pipeline(options=options)
+  pipeline(p)
+  result = p.run()
+  result.wait_until_finish()
+  for counter in result.metrics().query()['counters']:
+    logging.info('%s counter: %s.', counter.key.metric.name, counter)
+  logging.info('Completing beam pipeline.')
+
+
+def beam_convert_tfexamples(
+    root: beam.Pipeline,
+    data_list: Sequence[Dict[str, Any]],
+    convert_fn: Callable[[Dict[str, Any]], tf.train.Example],
+    output_dir: str,
+    num_shards: Sequence[int],
+) -> None:
+  """Constructs beam pipelines to convert train, val, test TF Examples."""
+  names = [TRAIN_TFRECORD_NAME, VALIDATION_TFRECORD_NAME, TEST_TFRECORD_NAME]
+  split_data = (
+      root
+      | 'Create PCollection' >> beam.Create(data_list)
+      | 'Data split' >> beam.Partition(partition_by_ml_use, 3)
+  )
+  for i in range(3):
+    ml_use: str = ALL_ML_USES[i]
+    num_shard = num_shards[i]
+    output_prefix = os.path.join(output_dir, names[i])
+    _ = (
+        split_data[i]
+        | f'Convert {ml_use} TF Examples'
+        >> beam.ParDo(WriteToTFRecord(output_prefix, num_shard, convert_fn))
+        | f'Group {ml_use} TF Record files' >> beam.GroupBy(lambda x: x[0])
+        | f'Merge {ml_use} TF Record files'
+        >> beam.Map(merge_tfrecords_func(output_prefix, num_shard))
+    )
+
+
+def merge_tfrecords_func(output_prefix: str, num_shard: int) -> ...:
+  """Returns a function to merge sharded worker output into expected shards."""
+  output_prefix = fileutils.force_gcs_fuse_path(output_prefix)
+
+  def merge_tfrecords(worker_output: Tuple[int, Sequence[Tuple[int, str]]]):
+    idx = worker_output[0]
+    files: Sequence[str] = np.unique([x[1] for x in worker_output[1]])
+    output_file = f'{output_prefix}-{idx:05d}-of-{num_shard:05d}'
+    with open(output_file, 'wb') as f:
+      for file in files:
+        logging.info('Merging %s.', file)
+        file = fileutils.force_gcs_fuse_path(file)
+        with open(file, 'rb') as fin:
+          while True:
+            data = fin.read(READ_CHUNK_SIZE)
+            if not data:
+              break
+            f.write(data)
+        os.remove(file)
+
+  return merge_tfrecords

community-content/vertex_vision_model_garden/model_oss/data_converter/convert_coco_label_map_iod.py

@@ -0,0 +1,111 @@
+r"""Converts COCO labels as yamls for model garden playground (IOD).
+"""
+
+import os
+import urllib.request
+
+from absl import app
+from absl import flags
+import tensorflow as tf
+import yaml
+
+from object_detection.utils import label_map_util
+
+_CONVERT_LABEL_TYPE_COCO_80 = 'coco_80'
+_CONVERT_LABEL_TYPE_COCO_91 = 'coco_91'
+
+_CONVERT_LABEL_TYPE = flags.DEFINE_enum(
+    'convert_label_type',
+    None,
+    [
+        _CONVERT_LABEL_TYPE_COCO_80,
+        _CONVERT_LABEL_TYPE_COCO_91,
+    ],
+    'Different types of label type conversion.',
+    required=True,
+)
+
+_TEMPORARY_PATH = flags.DEFINE_string(
+    'temporary_path',
+    None,
+    'The tempory path.',
+    required=True,
+)
+
+_OUTPUT_YAML_FILEPATH = flags.DEFINE_string(
+    'output_yaml_filepath',
+    None,
+    'The output yaml filepath.',
+    required=True,
+)
+
+
+def convert_coco_label_map_91(
+    output_yaml_filepath: str,
+) -> None:
+  """Converts coco label map 91."""
+  input_proto_filepath = 'https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/data/mscoco_label_map.pbtxt'
+  local_input_proto_filepath = os.path.join(
+      _TEMPORARY_PATH.value, 'mscoco_label_map.pbtxt'
+  )
+  with open(local_input_proto_filepath, 'w') as writer:
+    contents = (
+        urllib.request.urlopen(input_proto_filepath).read().decode('utf-8')
+    )
+    writer.write(contents)
+
+  label_map = label_map_util.load_labelmap(local_input_proto_filepath)
+  label_map_dict = label_map_util.get_label_map_dict(
+      label_map, use_display_name=True
+  )
+  swapped_label_map_dict = {v: k for k, v in label_map_dict.items()}
+  print(swapped_label_map_dict)
+
+  # Saves new label maps as yamls.
+  with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer:
+    writer.write(yaml.dump(swapped_label_map_dict))
+
+
+def convert_coco_label_map_80(
+    output_yaml_filepath: str,
+) -> None:
+  """Converts coco label map 80."""
+  # Loads label maps from texts.
+  input_text_filepath = 'https://gist.githubusercontent.com/AruniRC/7b3dadd004da04c80198557db5da4bda/raw/2f10965ace1e36c4a9dca76ead19b744f5eb7e88/ms_coco_classnames.txt'
+  local_input_text_filepath = os.path.join(
+      _TEMPORARY_PATH.value, 'ms_coco_classnames.txt'
+  )
+  with open(local_input_text_filepath, 'w') as writer:
+    contents = (
+        urllib.request.urlopen(input_text_filepath).read().decode('utf-8')
+    )
+    writer.write(contents)
+  with open(local_input_text_filepath, 'r') as file:
+    content = file.read()
+    label_map = yaml.safe_load(content)
+
+  # Removes background in label maps.
+  new_label_map = {}
+  for k, v in label_map.items():
+    if k == 0:
+      continue
+    new_label_map[k - 1] = v
+  print(new_label_map)
+  # Saves new label maps as yamls.
+  with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer:
+    writer.write(yaml.dump(new_label_map))
+
+
+def main(_) -> None:
+  if _CONVERT_LABEL_TYPE.value == _CONVERT_LABEL_TYPE_COCO_80:
+    convert_coco_label_map_80(_OUTPUT_YAML_FILEPATH.value)
+  elif _CONVERT_LABEL_TYPE.value == _CONVERT_LABEL_TYPE_COCO_91:
+    convert_coco_label_map_91(
+        _OUTPUT_YAML_FILEPATH.value,
+    )
+  else:
+    print('Not supported convert label type: ', _CONVERT_LABEL_TYPE.value)
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/data_converter/convert_imagenet_label_map.py

@@ -0,0 +1,86 @@
+r"""Converts ImageNet label texts as yamls for model garden playground.
+
+# ImageNet1K will have label maps with background.
+"""
+
+import urllib.request
+from absl import app
+from absl import flags
+import tensorflow as tf
+import yaml
+
+
+_INPUT_TEXT_FILEPATH = flags.DEFINE_string(
+    'input_text_filepath',
+    None,
+    'The input text filepath.',
+    required=True,
+)
+
+_ADD_BACKGROUND_LABEL = flags.DEFINE_boolean(
+    'add_background_label',
+    None,
+    'Whether or not add background labels.',
+    required=True,
+)
+
+_ADD_IDS = flags.DEFINE_boolean(
+    'add_ids',
+    None,
+    'Whether or not add ids.',
+    required=True,
+)
+
+_OUTPUT_YAML_FILEPATH = flags.DEFINE_string(
+    'output_yaml_filepath',
+    None,
+    'The output yaml filepath.',
+    required=True,
+)
+
+
+def convert_imagenet_label_map_from_text_to_yaml(
+    input_text_filepath: str,
+    add_background_label: bool,
+    add_ids: bool,
+    output_yaml_filepath: str,
+) -> None:
+  """Converts imagenet label map from text to yamls."""
+  label_map = {}
+
+  # Shifts all keys by 1, and add 0 as 'background'.
+  if add_background_label:
+    label_map = yaml.safe_load(
+        urllib.request.urlopen(input_text_filepath).read()
+    )
+    new_label_map = {}
+    for key, value in label_map.items():
+      new_label_map[key + 1] = value
+    new_label_map[0] = 'background'
+    label_map = new_label_map
+
+  # Adds maps from id to each line.
+  if add_ids:
+    lines = urllib.request.urlopen(input_text_filepath).readlines()
+    current_id = 0
+    for line in lines:
+      label_map[current_id] = line.decode('ascii').strip()
+      print(label_map[current_id])
+      current_id += 1
+
+  # Saves new label maps as yamls.
+  with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer:
+    writer.write(yaml.dump(label_map))
+
+
+def main(_) -> None:
+  convert_imagenet_label_map_from_text_to_yaml(
+      _INPUT_TEXT_FILEPATH.value,
+      _ADD_BACKGROUND_LABEL.value,
+      _ADD_IDS.value,
+      _OUTPUT_YAML_FILEPATH.value,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_icn_lib.py

@@ -0,0 +1,199 @@
+"""Converts ICN CSV/JSONL files to TFRecord with apache beam."""
+
+import json
+from os import path
+from typing import Any, Dict, Sequence, Union, cast
+
+from absl import logging
+import apache_beam as beam
+import pandas as pd
+import tensorflow as tf
+
+from data_converter import common_lib
+
+
+_COLUMN_NAMES = [
+    common_lib.COLUMN_NAME_ML_USE,
+    common_lib.COLUMN_NAME_GCS_FILE_PATH,
+    common_lib.COLUMN_NAME_LABEL,
+]
+_JSON_GCS_URI_KEY = 'imageGcsUri'
+_JSON_CLASS_ANNOTATION_KEY = 'classificationAnnotation'
+_JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels'
+_JSON_CLASS_NAME_KEY = 'displayName'
+_JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use'
+
+
+def build_tf_example(element: Dict[str, Union[str, int]]) -> tf.train.Example:
+  """Builds a TF Example from an image uri and label.
+
+  Args:
+    element: A dict with the keys gcs_file_path and label.
+
+  Returns:
+    The created TF Example.
+  """
+  image_uri = cast(str, element[common_lib.COLUMN_NAME_GCS_FILE_PATH])
+  label = cast(int, element[common_lib.COLUMN_NAME_LABEL])
+  image_bytes, shape = common_lib.encode_image(image_uri, image_format='jpeg')
+  features = tf.train.Features(
+      feature={
+          'image/encoded': common_lib.convert_to_feature(image_bytes),
+          'image/format': common_lib.convert_to_string_feature('jpeg'),
+          'image/height': common_lib.convert_to_feature(shape[0]),
+          'image/width': common_lib.convert_to_feature(shape[1]),
+          'image/class/label': common_lib.convert_to_feature(label),
+      },
+  )
+  return tf.train.Example(features=features)
+
+
+def _run_convert_pipeline(
+    output_dir: str, df: pd.DataFrame, num_shards: Sequence[int]
+) -> None:
+  """Starts a Beam pipeline to write DataFrame as TF Records.
+
+  Args:
+    output_dir: TF Records output directory.
+    df: DataFrame to convert from.
+    num_shards: Number of shards for train/validation/test TFRecord files.
+  """
+  images_list = df.to_dict('records')
+
+  def pipeline(root: beam.Pipeline):
+    common_lib.beam_convert_tfexamples(
+        root,
+        images_list,
+        build_tf_example,
+        output_dir,
+        num_shards,
+    )
+
+  common_lib.run_beam_pipeline(pipeline)
+
+
+def _convert_df_to_tfrecord(
+    df: pd.DataFrame,
+    output_dir: str,
+    split_ratio: Sequence[float],
+    num_shard: Sequence[int],
+) -> None:
+  """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir.
+
+  Args:
+    df: DataFrame to convert.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  common_lib.replace_unassigned_ml_use(
+      df[common_lib.COLUMN_NAME_ML_USE], split_ratio
+  )
+
+  # Converts labels to integers as required by training.
+  new_labels, label_map = common_lib.create_label_map(
+      df[common_lib.COLUMN_NAME_LABEL]
+  )
+  df[common_lib.COLUMN_NAME_LABEL] = new_labels
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  _run_convert_pipeline(output_dir, df, num_shard)
+
+
+def convert_csv_to_tfrecord(
+    input_csv: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The csv format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#csv.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_csv: Name of the csv file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_csv, 'r') as f:
+    df: pd.DataFrame = pd.read_csv(
+        f, header=None, names=_COLUMN_NAMES, on_bad_lines='warn'
+    )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)
+
+
+def convert_jsonl_to_tfrecord(
+    input_jsonl: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The JSONL format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#json-lines.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_jsonl: Name of the JSONL file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  df_rows = []
+  with tf.io.gfile.GFile(input_jsonl, 'r') as f:
+    lines = f.read().rstrip().splitlines()
+
+  for i, line in enumerate(lines, 1):
+    try:
+      item: Dict[str, Any] = json.loads(line)
+
+      gcs_uri = item.get(_JSON_GCS_URI_KEY)
+      label = item.get(_JSON_CLASS_ANNOTATION_KEY, {}).get(_JSON_CLASS_NAME_KEY)
+      if not gcs_uri or not label:
+        logging.warning('Invalid JSON at line %d, skipped.', i)
+        continue
+
+      ml_use = item.get(_JSON_RESOURCE_LABEL_KEY, {}).get(
+          _JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED
+      )
+    except (json.JSONDecodeError, AttributeError):
+      logging.warning('Invalid JSON at line %d, skipped.', i)
+      continue
+
+    df_rows.append([ml_use, gcs_uri, label])
+
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=[
+          common_lib.COLUMN_NAME_ML_USE,
+          common_lib.COLUMN_NAME_GCS_FILE_PATH,
+          common_lib.COLUMN_NAME_LABEL,
+      ],
+  )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_iod_lib.py

@@ -0,0 +1,430 @@
+"""Converts IOD dataset files to TFRecord with apache beam."""
+
+import collections
+import json
+from os import path
+from typing import Any, Dict, Sequence
+
+from absl import logging
+import apache_beam as beam
+import pandas as pd
+import tensorflow as tf
+
+from data_converter import common_lib
+from util import constants
+
+COLUMN_NAME_LABEL_INT = 'label_int'
+_COLUMN_NAME_XMIN = 'X_MIN'
+_COLUMN_NAME_YMIN = 'Y_MIN'
+_COLUMN_NAME_XMAX = 'X_MAX'
+_COLUMN_NAME_YMAX = 'Y_MAX'
+COLUMN_NAMES = [
+    common_lib.COLUMN_NAME_ML_USE,
+    common_lib.COLUMN_NAME_GCS_FILE_PATH,
+    common_lib.COLUMN_NAME_LABEL,
+    _COLUMN_NAME_XMIN,
+    _COLUMN_NAME_YMIN,
+    'XMAX_NOT_USED',
+    'YMIN_NOT_USED',
+    _COLUMN_NAME_XMAX,
+    _COLUMN_NAME_YMAX,
+    'XMIN_NOT_USED',
+    'YMAX_NOT_USED',
+]
+_BOUNDING_BOX_COLUMNS = [
+    _COLUMN_NAME_XMIN,
+    _COLUMN_NAME_YMIN,
+    _COLUMN_NAME_XMAX,
+    _COLUMN_NAME_YMAX,
+]
+_JSON_BBOX_ANNOTATIONS_KEY = 'boundingBoxAnnotations'
+_JSON_DISPLAY_NAME_KEY = 'displayName'
+_JSON_X_MIN_KEY = 'xMin'
+_JSON_X_MAX_KEY = 'xMax'
+_JSON_Y_MIN_KEY = 'yMin'
+_JSON_Y_MAX_KEY = 'yMax'
+
+
+def build_tf_example(image_row: Dict[str, Any]) -> tf.train.Example:
+  """Builds a TF Example from an image row.
+
+  Args:
+    image_row: A dictionary containing information about the image, such as its
+      GCS uri, labels, and bounding box coordinates.
+
+  Returns:
+    A tf.train.Example containing the encoded image and optionally a
+      bounding box and label.
+  """
+  image_uri = image_row[common_lib.COLUMN_NAME_GCS_FILE_PATH]
+  image_bytes, shape = common_lib.encode_image(image_uri, image_format='jpeg')
+  feature = {
+      'image/encoded': common_lib.convert_to_feature(image_bytes),
+      'image/format': common_lib.convert_to_string_feature('jpeg'),
+      'image/height': common_lib.convert_to_feature(shape[0]),
+      'image/width': common_lib.convert_to_feature(shape[1]),
+      'image/source_id': common_lib.convert_to_string_feature(image_uri),
+      'image/object/bbox/xmin': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_XMIN]
+      ),
+      'image/object/bbox/ymin': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_YMIN]
+      ),
+      'image/object/bbox/xmax': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_XMAX]
+      ),
+      'image/object/bbox/ymax': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_YMAX]
+      ),
+      'image/object/class/text': common_lib.convert_to_list_string_feature(
+          image_row[common_lib.COLUMN_NAME_LABEL]
+      ),
+      'image/object/class/label': common_lib.convert_to_feature(
+          image_row[COLUMN_NAME_LABEL_INT]
+      ),
+  }
+  return tf.train.Example(features=tf.train.Features(feature=feature))
+
+
+def _run_convert_pipeline(
+    output_dir: str,
+    image_rows: Sequence[Dict[str, Any]],
+    num_shards: Sequence[int],
+) -> None:
+  """Starts a Beam pipeline to write DataFrame as TF Records.
+
+  Args:
+    output_dir: TF Records output directory.
+    image_rows: Contains all necessary information to create a TF Example.
+    num_shards: Number of shards for train/validation/test TFRecord files.
+  """
+
+  def pipeline(root: beam.Pipeline):
+    common_lib.beam_convert_tfexamples(
+        root,
+        image_rows,
+        build_tf_example,
+        output_dir,
+        num_shards,
+    )
+
+  common_lib.run_beam_pipeline(pipeline)
+
+
+def _convert_df_to_tfrecord(
+    df: pd.DataFrame,
+    output_dir: str,
+    split_ratio: Sequence[float],
+    num_shard: Sequence[int],
+) -> None:
+  """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir.
+
+  Args:
+    df: DataFrame to convert.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Specify bounding box columns to be numeric.
+  df[_BOUNDING_BOX_COLUMNS] = df[_BOUNDING_BOX_COLUMNS].apply(pd.to_numeric)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  dropped_row_num += drop_rows_without_bbox(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  # Converts labels to integers as required by training.
+  int_labels, label_map = common_lib.create_label_map(
+      df[common_lib.COLUMN_NAME_LABEL]
+  )
+  df[COLUMN_NAME_LABEL_INT] = int_labels
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  image_rows = _condense_bounding_boxes(df.to_dict(orient='records'))
+  ml_uses = [row[common_lib.COLUMN_NAME_ML_USE] for row in image_rows]
+  common_lib.replace_unassigned_ml_use(ml_uses, split_ratio)
+  common_lib.merge_seq_into_dicts(
+      common_lib.COLUMN_NAME_ML_USE, ml_uses, image_rows
+  )
+
+  _run_convert_pipeline(output_dir, image_rows, num_shard)
+
+
+def _condense_bounding_boxes(
+    image_rows: Sequence[Dict[str, Any]]
+) -> Sequence[Dict[str, Any]]:
+  """Gather all the bounding boxes in an image and put them in the same dictionary.
+
+  Args:
+    image_rows: List of dictionaries, each containing information about the
+      image, such as its GCS uri, labels, and bounding box coordinates.
+
+  Returns:
+    List of dictionaries such that each contains all the bounding boxes for a
+    given gcs_file_path.
+
+  Raises:
+    RuntimeError: This is raised when the input data contains images that have
+      annotations in different ml_use classes.
+  """
+  output = {}
+  for image_row in image_rows:
+    ml_use = image_row[common_lib.COLUMN_NAME_ML_USE]
+    gcs_file_path = image_row[common_lib.COLUMN_NAME_GCS_FILE_PATH]
+    label = image_row[common_lib.COLUMN_NAME_LABEL]
+    xmin = image_row[_COLUMN_NAME_XMIN]
+    ymin = image_row[_COLUMN_NAME_YMIN]
+    xmax = image_row[_COLUMN_NAME_XMAX]
+    ymax = image_row[_COLUMN_NAME_YMAX]
+    label_int = image_row[COLUMN_NAME_LABEL_INT]
+    if gcs_file_path in output:
+      d = output[gcs_file_path]
+      if ml_use != common_lib.ML_USE_UNASSIGNED:
+        if d[common_lib.COLUMN_NAME_ML_USE] == common_lib.ML_USE_UNASSIGNED:
+          d[common_lib.COLUMN_NAME_ML_USE] = ml_use
+        elif ml_use != d[common_lib.COLUMN_NAME_ML_USE]:
+          raise RuntimeError(
+              f'Image {gcs_file_path} can only be placed in one of'
+              f' training/validation/test. It is currently in {ml_use} and'
+              f' {d[common_lib.COLUMN_NAME_ML_USE]}.'
+          )
+      d[common_lib.COLUMN_NAME_LABEL].append(label)
+      d[_COLUMN_NAME_XMIN].append(xmin)
+      d[_COLUMN_NAME_YMIN].append(ymin)
+      d[_COLUMN_NAME_XMAX].append(xmax)
+      d[_COLUMN_NAME_YMAX].append(ymax)
+      d[COLUMN_NAME_LABEL_INT].append(label_int)
+    else:
+      output[gcs_file_path] = {
+          common_lib.COLUMN_NAME_ML_USE: ml_use,
+          common_lib.COLUMN_NAME_GCS_FILE_PATH: gcs_file_path,
+          common_lib.COLUMN_NAME_LABEL: [label],
+          _COLUMN_NAME_XMIN: [xmin],
+          _COLUMN_NAME_YMIN: [ymin],
+          _COLUMN_NAME_XMAX: [xmax],
+          _COLUMN_NAME_YMAX: [ymax],
+          COLUMN_NAME_LABEL_INT: [label_int],
+      }
+  return list(output.values())
+
+
+def convert_csv_to_tfrecord(
+    input_csv: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The csv format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#csv.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_csv: Name of the csv file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the train, validation, and test splits for
+      unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_csv, 'r') as f:
+    df: pd.DataFrame = pd.read_csv(
+        f, header=None, names=COLUMN_NAMES, on_bad_lines='warn'
+    )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)
+
+
+def drop_rows_without_bbox(df: pd.DataFrame) -> int:
+  """Drops DataFrame rows without bounding_boxes.
+
+  Args:
+    df: The DataFrame to process in place.
+
+  Returns:
+    The number of rows dropped.
+  """
+  invalid_rows = df.index[~(df[_BOUNDING_BOX_COLUMNS].notnull().all(axis=1))]
+  dropped_num = len(invalid_rows)
+  if dropped_num > 0:
+    invalid_df = df.loc[invalid_rows].to_dict(orient='records')
+    for entry in invalid_df:
+      logging.warning('Skipping entry due to missing bounding box: %s.', entry)
+    df.drop(invalid_rows, inplace=True)
+    df.reset_index(drop=True, inplace=True)
+  return dropped_num
+
+
+def convert_coco_json_categories_to_label_map(
+    categories: Sequence[Dict[str, Any]]
+) -> Dict[int, str]:
+  return {category['id']: category['name'] for category in categories}
+
+
+def convert_coco_json_to_tfrecord(
+    input_coco_json: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The COCO json format is shown here: https://cocodataset.org/#format-data.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_coco_json: Name of coco json file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the train, validation, and test splits for
+      dataset.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_coco_json, 'r') as f:
+    coco_json = json.load(f)
+  # Writes label map from coco json categories.
+  label_map = convert_coco_json_categories_to_label_map(
+      coco_json[constants.COCO_JSON_CATEGORIES]
+  )
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writes label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  img_to_anns = collections.defaultdict(list)
+  imgs = {}
+  if constants.COCO_JSON_ANNOTATIONS in coco_json:
+    for ann in coco_json[constants.COCO_JSON_ANNOTATIONS]:
+      img_to_anns[ann[constants.COCO_JSON_ANNOTATION_IMAGE_ID]].append(ann)
+
+  if constants.COCO_JSON_IMAGES in coco_json:
+    for img in coco_json[constants.COCO_JSON_IMAGES]:
+      imgs[img[constants.COCO_JSON_IMAGE_ID]] = img
+
+  df_rows = []
+
+  for image_id, annotations in img_to_anns.items():
+    img = imgs[image_id]
+    for ann in annotations:
+      xmin, ymin, xmax, ymax = common_lib.reformat_bbox(
+          ann[constants.COCO_ANNOTATION_BBOX],
+          img[constants.COCO_JSON_IMAGE_WIDTH],
+          img[constants.COCO_JSON_IMAGE_HEIGHT],
+      )
+      df_rows.append([
+          common_lib.ML_USE_UNASSIGNED,
+          img[constants.COCO_JSON_IMAGE_COCO_URL],
+          label_map[ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID]],
+          xmin,
+          ymin,
+          xmax,
+          ymin,
+          xmax,
+          ymax,
+          xmin,
+          ymax,
+          ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID],
+      ])
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=COLUMN_NAMES + [COLUMN_NAME_LABEL_INT],
+  )
+
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Species bounding box columns to be numeric.
+  df[_BOUNDING_BOX_COLUMNS] = df[_BOUNDING_BOX_COLUMNS].apply(pd.to_numeric)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  dropped_row_num += drop_rows_without_bbox(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  image_rows = _condense_bounding_boxes(df.to_dict(orient='records'))
+  ml_uses = [row[common_lib.COLUMN_NAME_ML_USE] for row in image_rows]
+  common_lib.replace_unassigned_ml_use(ml_uses, split_ratio)
+  common_lib.merge_seq_into_dicts(
+      common_lib.COLUMN_NAME_ML_USE, ml_uses, image_rows
+  )
+
+  _run_convert_pipeline(output_dir, image_rows, num_shard)
+
+
+def convert_jsonl_to_tfrecord(
+    input_jsonl: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The JSONL format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#json-lines.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_jsonl: Name of the JSONL file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  df_rows = []
+  with tf.io.gfile.GFile(input_jsonl, 'r') as f:
+    lines = f.read().rstrip().splitlines()
+
+  for i, line in enumerate(lines, start=1):
+    try:
+      item: Dict[str, Any] = json.loads(line)
+    except (json.JSONDecodeError, AttributeError):
+      logging.warning('Invalid JSON at line %d skipped.', i)
+      continue
+
+    gcs_uri = item.get(common_lib.JSON_GCS_URI_KEY)
+    if not gcs_uri:
+      logging.warning(
+          'Invalid JSON at line %d skipped. Missing gcs_uri_key.', i
+      )
+      continue
+    ml_use = item.get(common_lib.JSON_RESOURCE_LABEL_KEY, {}).get(
+        common_lib.JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED
+    )
+
+    for bbox in item.get(_JSON_BBOX_ANNOTATIONS_KEY, []):
+      label = bbox.get(_JSON_DISPLAY_NAME_KEY)
+      xmin = bbox.get(_JSON_X_MIN_KEY)
+      ymin = bbox.get(_JSON_Y_MIN_KEY)
+      xmax = bbox.get(_JSON_X_MAX_KEY)
+      ymax = bbox.get(_JSON_Y_MAX_KEY)
+
+      df_rows.append([ml_use, gcs_uri, label, xmin, ymin, xmax, ymax])
+
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=[
+          common_lib.COLUMN_NAME_ML_USE,
+          common_lib.COLUMN_NAME_GCS_FILE_PATH,
+          common_lib.COLUMN_NAME_LABEL,
+          _COLUMN_NAME_XMIN,
+          _COLUMN_NAME_YMIN,
+          _COLUMN_NAME_XMAX,
+          _COLUMN_NAME_YMAX,
+      ],
+  )
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_isg_lib.py

@@ -0,0 +1,328 @@
+"""Python script to convert different file formats for ISG to tfrecords."""
+
+import hashlib
+import os
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+
+from absl import logging
+import apache_beam as beam
+from apache_beam.io import tfrecordio
+import cv2
+import numpy as np
+from pycocotools import coco
+import tensorflow as tf
+import yaml
+
+from data_converter import common_lib
+from util import constants
+from util import fileutils
+
+_IMAGE_FORMAT = 'PNG'
+
+
+def build_tf_example(
+    image_info: dict[str, Union[str, int]],
+    segmentation_image: List[List[int]],
+    output_shape: Optional[Tuple[int, int]] = None,
+) -> tf.train.Example:
+  """Encodes an image and its segmentation mask into a tf.train.Example.
+
+  Args:
+    image_info: A dictionary containing information about the image, such as its
+      file name, height, and width.
+    segmentation_image: 2D image in list of lists having category ids.
+    output_shape: The desired output shape of the image. If None, the original
+      image shape will be used.
+
+  Returns:
+    A tf.train.Example containing the encoded image and segmentation mask.
+
+  Raises:
+    IOError: If image cannot be found in the path.
+  """
+  file_name = image_info[constants.COCO_JSON_FILE_NAME]
+  height = int(image_info[constants.COCO_JSON_IMAGE_HEIGHT])
+  width = int(image_info[constants.COCO_JSON_IMAGE_WIDTH])
+
+  segmentation_image = np.expand_dims(
+      np.asarray(segmentation_image, dtype=np.int32), axis=-1
+  )
+  _, encoded_seg = cv2.imencode(f'.{_IMAGE_FORMAT.lower()}', segmentation_image)
+  encoded_seg = encoded_seg.tobytes()
+
+  encoded_img, _ = common_lib.encode_image(
+      image_info[constants.COCO_JSON_IMAGE_COCO_URL],
+      output_shape=output_shape,
+      image_format=_IMAGE_FORMAT.lower(),
+  )
+
+  key = hashlib.sha256(encoded_img).hexdigest()
+
+  return tf.train.Example(
+      features=tf.train.Features(
+          feature={
+              'image/height': common_lib.convert_to_feature(height),
+              'image/width': common_lib.convert_to_feature(width),
+              'image/filename': common_lib.convert_to_string_feature(file_name),
+              'image/sha256': common_lib.convert_to_string_feature(key),
+              'image/encoded': common_lib.convert_to_feature(encoded_img),
+              'image/format': common_lib.convert_to_string_feature(
+                  _IMAGE_FORMAT
+              ),
+              'image/segmentation/class/encoded': common_lib.convert_to_feature(
+                  encoded_seg
+              ),
+              'image/segmentation/class/format': (
+                  common_lib.convert_to_string_feature(_IMAGE_FORMAT)
+              ),
+              'image/segmentation/class/height': common_lib.convert_to_feature(
+                  height
+              ),
+              'image/segmentation/class/width': common_lib.convert_to_feature(
+                  width
+              ),
+          }
+      )
+  )
+
+
+class AcquireTFExampleDoFn(beam.DoFn):
+  """Beam DoFn to build TF Examples from a single row of image_info data."""
+
+  # These tags will be used to tag the outputs of this DoFn.
+  output_tag_train = constants.ML_USE_TRAINING
+  output_tag_validation = constants.ML_USE_VALIDATION
+  output_tag_test = constants.ML_USE_TEST
+
+  valid_ml_use_set = set(
+      [output_tag_train, output_tag_validation, output_tag_test]
+  )
+
+  def __init__(self, output_shape: Optional[Tuple[int, int]] = None):
+    self.acquired_examples_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Success'
+    )
+    self.failure_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Failure'
+    )
+    self.output_shape = output_shape
+
+  def process(
+      self,
+      row: Tuple[str, Dict[str, Union[str, int]], List[List[int]]],
+  ) -> Iterator[tf.train.Example]:
+    ml_use, image_info, annotation_info = row
+    if ml_use not in self.valid_ml_use_set:
+      logging.warning('ml_use invalid: %s', ml_use)
+      self.failure_counter.inc()
+      return
+
+    try:
+      tf_example = build_tf_example(
+          image_info, annotation_info, self.output_shape
+      )
+    except IOError as e:
+      logging.warning('Failed to build TF Example: %s', e)
+      self.failure_counter.inc()
+    else:
+      self.acquired_examples_counter.inc()
+      yield beam.pvalue.TaggedOutput(ml_use, tf_example)
+
+
+def _define_data_conversion_pipeline(
+    root: beam.Pipeline,
+    ml_use_rows: List[str],
+    image_rows: List[Dict[str, Union[str, int]]],
+    segmentation_rows: List[List[List[int]]],
+    output_dir: str,
+    output_shape: Optional[Tuple[int, int]],
+    num_shard_list: List[int],
+):
+  """Define a data conversion pipeline.
+
+  Args:
+    root: A Beam pipeline.
+    ml_use_rows: List containing the ml_use.
+    image_rows: List of dictionaries containing information about the image,
+      such as its file name, height, and width.
+    segmentation_rows: List of 2D images of integers representing segmentation
+      masks.
+    output_dir: Directory where the output TFRecords will be written.
+    output_shape: Desired output shape of the image. If None, the original image
+      shape will be used.
+    num_shard_list: Number of shards to write to each output TFRecord.
+
+  Returns:
+    A Beam pipeline.
+  """
+  train, validation, test = (
+      root
+      | 'Load ml use and image rows to beam'
+      >> beam.Create(zip(ml_use_rows, image_rows, segmentation_rows))
+      | 'Build TF Examples'
+      >> beam.ParDo(AcquireTFExampleDoFn(output_shape)).with_outputs(
+          AcquireTFExampleDoFn.output_tag_train,
+          AcquireTFExampleDoFn.output_tag_validation,
+          AcquireTFExampleDoFn.output_tag_test,
+      )
+  )
+
+  # Save each split to TFRecord.
+  _ = train | 'Save train split to TFRecord' >> tfrecordio.WriteToTFRecord(
+      os.path.join(output_dir, common_lib.TRAIN_TFRECORD_NAME),
+      coder=beam.coders.ProtoCoder(tf.train.Example),
+      num_shards=num_shard_list[0],
+  )
+  _ = (
+      validation
+      | 'Save validation split to TFRecord'
+      >> tfrecordio.WriteToTFRecord(
+          os.path.join(output_dir, common_lib.VALIDATION_TFRECORD_NAME),
+          coder=beam.coders.ProtoCoder(tf.train.Example),
+          num_shards=num_shard_list[1],
+      )
+  )
+  _ = test | 'Save test split to TFRecord' >> tfrecordio.WriteToTFRecord(
+      os.path.join(output_dir, common_lib.TEST_TFRECORD_NAME),
+      coder=beam.coders.ProtoCoder(tf.train.Example),
+      num_shards=num_shard_list[2],
+  )
+
+
+def _image_info_to_segmentation_image(
+    img: Dict[str, Any],
+    coco_dataset: coco.COCO,
+    label_id_by_category_id: Dict[int, int],
+) -> List[List[int]]:
+  """Convert image information to a segmentation image.
+
+  Args:
+      img: The image information.
+      coco_dataset: The COCO dataset.
+      label_id_by_category_id: The mapping from label id used for training to
+        category_id defined in dataset.
+
+  Returns:
+      The segmentation image.
+
+  Raises:
+    ValueError: If the mask size does not match the image or if a pixel has
+      multiple labels.
+  """
+  seg_img = np.zeros(
+      shape=(
+          img[constants.COCO_JSON_IMAGE_HEIGHT],
+          img[constants.COCO_JSON_IMAGE_WIDTH],
+      ),
+      dtype=np.int32,
+  )
+  for ann in coco_dataset.imgToAnns[img[constants.COCO_JSON_IMAGE_ID]]:
+    new_category_id = ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID]
+    binary_mask = coco_dataset.annToMask(ann)
+    if seg_img.shape != binary_mask.shape:
+      raise ValueError(
+          'Binary mask does not have the same shape as image. image_id:'
+          f' {img["id"]}'
+      )
+    boolean_mask = binary_mask == 1
+    if (seg_img[boolean_mask] != 0).any():
+      raise ValueError(
+          'Error: Some pixels have more than one label in image_id:'
+          f' {img["id"]}.'
+      )
+    seg_img[boolean_mask] = label_id_by_category_id[new_category_id]
+
+  return seg_img.tolist()
+
+
+def get_input_rows(
+    coco_dataset: coco.COCO,
+    split_ratio: List[float],
+    label_id_by_category_id: Dict[int, int],
+) -> Tuple[List[str], List[Dict[str, Union[str, int]]], List[List[List[int]]]]:
+  """Get input rows for training and validation.
+
+  Args:
+      coco_dataset: The COCO dataset.
+      split_ratio: The split ratio for training and validation.
+      label_id_by_category_id: The mapping from label id used for training to
+        category_id defined in dataset.
+
+  Returns:
+      - A list of ml_use strings.
+      - A list of image informations.
+      - A list of segmentation images for the corresponding images.
+  """
+  image_rows = coco_dataset.dataset[constants.COCO_JSON_IMAGES]
+
+  segmentation_rows = [
+      _image_info_to_segmentation_image(
+          img, coco_dataset, label_id_by_category_id
+      )
+      for img in image_rows
+  ]
+
+  ml_use_rows = common_lib.create_ml_use_array_with_split(
+      len(image_rows), split_ratio
+  )
+  return ml_use_rows, image_rows, segmentation_rows
+
+
+def beam_build_tfrecord_from_coco_json(
+    input_json: str,
+    output_dir: str,
+    split_ratio: List[float],
+    num_shard_list: List[int],
+    output_shape: Optional[Tuple[int, int]] = None,
+) -> None:
+  """Builds TFRecord files from COCO dataset.
+
+  The output file names are `_TRAIN_TFRECORD_NAME`, `_VALIDATION_TFRECORD_NAME`,
+  and `_TEST_TFRECORD_NAME`.
+
+  Args:
+    input_json: Path to a COCO JSON or JSONL file.
+    output_dir: Directory to output the TFRecord files.
+    split_ratio: List of how to split entries to train, validation, and test
+      TFRecords.
+    num_shard_list: List of the number of shards for each TFRecord file.
+    output_shape: The desired output shape of the image. If None, the original
+      image shape will be used.
+  """
+  # `coco` cannot access gcs uri. Use gcsfuse, it is faster.
+  input_json = fileutils.force_gcs_fuse_path(input_json)
+  coco_dataset = coco.COCO(input_json)
+
+  label_map = {}
+  label_id_by_category_id = {}
+  for idx, category in enumerate(
+      coco_dataset.dataset[constants.COCO_JSON_CATEGORIES], start=1
+  ):
+    label_map[idx] = category[constants.COCO_JSON_CATEGORY_NAME]
+    label_id_by_category_id[category[constants.COCO_JSON_CATEGORY_ID]] = idx
+  label_map_path = os.path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  with tf.io.gfile.GFile(
+      os.path.join(output_dir, 'label_id_by_category_id.yaml'), 'w'
+  ) as f:
+    yaml.dump(label_id_by_category_id, f)
+
+  ml_use_rows, image_rows, segmentation_rows = get_input_rows(
+      coco_dataset, split_ratio, label_id_by_category_id
+  )
+
+  def pipeline(root):
+    _define_data_conversion_pipeline(
+        root,
+        ml_use_rows,
+        image_rows,
+        segmentation_rows,
+        output_dir,
+        output_shape,
+        num_shard_list,
+    )
+
+  logging.info('Beginning beam pipeline to acquire tfrecords.')
+  common_lib.run_beam_pipeline(pipeline)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_main.py

@@ -0,0 +1,166 @@
+r"""Python script to convert user input data to training docker format.
+
+
+Note: the training format is designed to be tfrecord as in the design doc.
+If there are training efficiency issues for pytorch algorithms, we will also
+support pytorch formats as well.
+"""
+
+from absl import app
+from absl import flags
+from absl import logging
+from data_converter import common_lib
+from data_converter import data_converter_icn_lib
+from data_converter import data_converter_iod_lib
+from data_converter import data_converter_isg_lib
+from data_converter import data_converter_vcn_lib
+from util import constants
+
+
+_INPUT_FILE_PATH = flags.DEFINE_string(
+    'input_file_path',
+    None,
+    'Input file path.',
+    required=True,
+)
+_INPUT_FILE_TYPE = flags.DEFINE_enum(
+    'input_file_type',
+    None,
+    [
+        constants.INPUT_FILE_TYPE_CSV,
+        constants.INPUT_FILE_TYPE_JSONL,
+        constants.INPUT_FILE_TYPE_COCO_JSON,
+    ],
+    'Input file type.',
+    required=True,
+)
+_OBJECTIVE = flags.DEFINE_enum(
+    'objective',
+    None,
+    [
+        constants.OBJECTIVE_IMAGE_CLASSIFICATION,
+        constants.OBJECTIVE_IMAGE_OBJECT_DETECTION,
+        constants.OBJECTIVE_IMAGE_SEGMENTATION,
+        constants.OBJECTIVE_VIDEO_CLASSIFICATION,
+    ],
+    'The objective of this training job.',
+    required=True,
+)
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'The output directory for converted data and label map files.',
+    required=True,
+)
+_SPLIT_RATIO = flags.DEFINE_list(
+    'split_ratio',
+    '0.8,0.1,0.1',
+    'Proportion of data to split into train/validation/test.',
+)
+_NUM_SHARD = flags.DEFINE_list(
+    'num_shard', '10,10,10', 'The number of shards for train/validation/test.'
+)
+_OUTPUT_FPS = flags.DEFINE_integer(
+    'output_fps', 5, 'For videos only. The output frames rate per second.'
+)
+
+
+def main(_) -> None:
+  logging.info(
+      (
+          'Start data converter on: %s (type: %s) with split: %s for %s'
+          ' (shard=%s), and output to %s.'
+      ),
+      _INPUT_FILE_PATH.value,
+      _INPUT_FILE_TYPE.value,
+      _SPLIT_RATIO.value,
+      _OBJECTIVE.value,
+      _NUM_SHARD.value,
+      _OUTPUT_DIR.value,
+  )
+  split_ratio = list(map(float, _SPLIT_RATIO.value))
+  num_shard = list(map(int, _NUM_SHARD.value))
+  common_lib.check_split_ratio(split_ratio)
+  common_lib.check_num_shard(num_shard)
+  if (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV
+  ):
+    data_converter_iod_lib.convert_csv_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL
+  ):
+    data_converter_iod_lib.convert_jsonl_to_tfrecord(
+        _INPUT_FILE_PATH.value, _OUTPUT_DIR.value, split_ratio, num_shard
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_COCO_JSON
+  ):
+    data_converter_iod_lib.convert_coco_json_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_SEGMENTATION:
+    data_converter_isg_lib.beam_build_tfrecord_from_coco_json(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV
+  ):
+    data_converter_icn_lib.convert_csv_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL
+  ):
+    data_converter_icn_lib.convert_jsonl_to_tfrecord(
+        _INPUT_FILE_PATH.value, _OUTPUT_DIR.value, split_ratio, num_shard
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_VIDEO_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV
+  ):
+    data_converter_vcn_lib.convert_csv_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        _OUTPUT_FPS.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_VIDEO_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL
+  ):
+    data_converter_vcn_lib.convert_jsonl_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        _OUTPUT_FPS.value,
+        split_ratio,
+        num_shard,
+    )
+  else:
+    raise NotImplementedError(
+        f'File format {_INPUT_FILE_TYPE.value} is not supported for'
+        f' {_OBJECTIVE.value}.'
+    )
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_vcn_lib.py

@@ -0,0 +1,289 @@
+"""Converts VCN CSV/JSONL files to TFRecord with apache beam."""
+
+import json
+from os import path
+from typing import Any, Dict, Iterator, Sequence, Union, cast
+
+from absl import logging
+import apache_beam as beam
+from apache_beam.io import tfrecordio
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+
+from data_converter import common_lib
+from util import constants
+
+
+_COLUMN_NAMES = [
+    common_lib.COLUMN_NAME_ML_USE,
+    common_lib.COLUMN_NAME_GCS_FILE_PATH,
+    common_lib.COLUMN_NAME_LABEL,
+    common_lib.COLUMN_NAME_START_SEC,
+    common_lib.COLUMN_NAME_END_SEC,
+]
+_JSON_GCS_URI_KEY = 'videoGcsUri'
+_JSON_CLASS_ANNOTATION_KEY = 'timeSegmentAnnotations'
+_JSON_CLASS_NAME_KEY = 'displayName'
+_JSON_START_TIME_KEY = 'startTime'
+_JSON_END_TIME_KEY = 'endTime'
+_JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels'
+_JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use'
+
+
+def build_tf_example(
+    video_uri: str,
+    label: int,
+    start_sec: float,
+    end_sec: float,
+    output_fps: int,
+) -> tf.train.SequenceExample:
+  """Builds a TF Example from a video clip.
+
+  Args:
+    video_uri: GCS URI to the video file.
+    label: Class label as an integer.
+    start_sec: Start timestamp of the video clip in seconds.
+    end_sec: End timestamp of the video clip in seconds.
+    output_fps: The output frame rate per second.
+
+  Returns:
+    The created TF Example.
+  """
+  frame_bytes = common_lib.encode_video(
+      video_uri, start_sec, end_sec, output_fps, image_format='jpg'
+  )
+  seq_example = tf.train.SequenceExample()
+  seq_example.context.feature['clip/label/index'].int64_list.value[:] = [label]
+  for frame in frame_bytes:
+    seq_example.feature_lists.feature_list.get_or_create(
+        'image/encoded'
+    ).feature.add().bytes_list.value[:] = [frame]
+
+  return seq_example
+
+
+class AcquireTFExampleDoFn(beam.DoFn):
+  """Beam DoFn to build TF Examples from a DataFrame row dict for VCN."""
+
+  def __init__(self, output_fps: int):
+    self._success_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Success'
+    )
+    self._failure_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Failure'
+    )
+    self._output_fps = output_fps
+
+  def process(
+      self, element: Dict[str, Union[float, int, str]]
+  ) -> Iterator[tf.train.SequenceExample]:
+    ml_use: str = cast(str, element[common_lib.COLUMN_NAME_ML_USE])
+    video_uri: str = cast(str, element[common_lib.COLUMN_NAME_GCS_FILE_PATH])
+
+    try:
+      label: int = int(element[common_lib.COLUMN_NAME_LABEL])
+      start_sec: float = float(element[common_lib.COLUMN_NAME_START_SEC])
+      end_sec: float = float(element[common_lib.COLUMN_NAME_END_SEC])
+
+      tf_example = build_tf_example(
+          video_uri,
+          label,
+          start_sec,
+          end_sec,
+          self._output_fps,
+      )
+      self._success_counter.inc()
+      yield beam.pvalue.TaggedOutput(ml_use, tf_example)
+    except (ValueError, IOError) as err:
+      logging.error('Failed to process %s', video_uri)
+      logging.exception(err)
+      self._failure_counter.inc()
+
+
+def _run_convert_pipeline(
+    output_dir: str,
+    df: pd.DataFrame,
+    num_shards: Sequence[int],
+    output_fps: int,
+) -> None:
+  """Starts a Beam pipeline to write DataFrame as TF Records.
+
+  Args:
+    output_dir: TF Records output directory.
+    df: DataFrame to convert from.
+    num_shards: Number of shards for train/validation/test TFRecord files.
+    output_fps: The output frame rate per second.
+  """
+  clip_list = df.to_dict('records')
+
+  def pipeline(root):
+    train, val, test = (
+        root
+        | 'Create PCollection' >> beam.Create(clip_list)
+        | 'Convert to TF Example'
+        >> beam.ParDo(AcquireTFExampleDoFn(output_fps)).with_outputs(
+            constants.ML_USE_TRAINING,
+            constants.ML_USE_VALIDATION,
+            constants.ML_USE_TEST,
+        )
+    )
+    _ = train | 'Save train TF Record' >> tfrecordio.WriteToTFRecord(
+        path.join(output_dir, common_lib.TRAIN_TFRECORD_NAME),
+        coder=beam.coders.ProtoCoder(tf.train.Example),
+        num_shards=num_shards[0],
+    )
+    _ = val | 'Save val TF Record' >> tfrecordio.WriteToTFRecord(
+        path.join(output_dir, common_lib.VALIDATION_TFRECORD_NAME),
+        coder=beam.coders.ProtoCoder(tf.train.Example),
+        num_shards=num_shards[1],
+    )
+    _ = test | 'Save test TF Record' >> tfrecordio.WriteToTFRecord(
+        path.join(output_dir, common_lib.TEST_TFRECORD_NAME),
+        coder=beam.coders.ProtoCoder(tf.train.Example),
+        num_shards=num_shards[2],
+    )
+
+  common_lib.run_beam_pipeline(pipeline)
+
+
+def _convert_df_to_tfrecord(
+    df: pd.DataFrame,
+    output_dir: str,
+    split_ratio: Sequence[float],
+    num_shard: Sequence[int],
+    output_fps: int,
+) -> None:
+  """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir.
+
+  Args:
+    df: DataFrame to convert.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+    output_fps: The output frame rate per second.
+  """
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  common_lib.replace_unassigned_ml_use(
+      df[common_lib.COLUMN_NAME_ML_USE], split_ratio
+  )
+
+  # Converts labels to integers as required by training.
+  new_labels, label_map = common_lib.create_label_map(
+      df[common_lib.COLUMN_NAME_LABEL]
+  )
+  df[common_lib.COLUMN_NAME_LABEL] = new_labels
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  # Missing start / end times are treated as 0, inf, respectively.
+  df[common_lib.COLUMN_NAME_START_SEC].fillna(0, inplace=True)
+  df[common_lib.COLUMN_NAME_END_SEC].fillna(np.inf, inplace=True)
+
+  _run_convert_pipeline(output_dir, df, num_shard, output_fps)
+
+
+def convert_csv_to_tfrecord(
+    input_csv: str,
+    output_dir: str,
+    output_fps: int,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The csv format is shown in
+    https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data#csv
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_csv: Name of the csv file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    output_fps: The output frame rate per second.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_csv, 'r') as f:
+    df: pd.DataFrame = pd.read_csv(
+        f, header=None, names=_COLUMN_NAMES, on_bad_lines='warn'
+    )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard, output_fps)
+
+
+def convert_jsonl_to_tfrecord(
+    input_jsonl: str,
+    output_dir: str,
+    output_fps: int,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The JSONL format is shown in
+    https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data#jsonl.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_jsonl: Name of the JSONL file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    output_fps: The output frame rate per second.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  df_rows = []
+  with tf.io.gfile.GFile(input_jsonl, 'r') as f:
+    lines = f.read().rstrip().splitlines()
+
+  for i, line in enumerate(lines, 1):
+    try:
+      item: Dict[str, Any] = json.loads(line)
+
+      gcs_uri = item.get(_JSON_GCS_URI_KEY)
+      if not gcs_uri:
+        logging.warning('Invalid JSON at line %d, skipped.', i)
+        continue
+
+      annotations = item.get(_JSON_CLASS_ANNOTATION_KEY, [])
+      ml_use = item.get(_JSON_RESOURCE_LABEL_KEY, {}).get(
+          _JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED
+      )
+
+      for j, annotation in enumerate(annotations):
+        label = annotation.get(_JSON_CLASS_NAME_KEY)
+        if not label:
+          logging.warning('Invalid annotation #%d at line %d, skipped.', j, i)
+          continue
+        # The example in external documentation uses strings like "1.0s", so we
+        # need to remove the "s" suffix.
+        start_time = annotation.get(_JSON_START_TIME_KEY, '0').removesuffix('s')
+        end_time = annotation.get(_JSON_END_TIME_KEY, 'inf').removesuffix('s')
+        df_rows.append([ml_use, gcs_uri, label, start_time, end_time])
+    except (json.JSONDecodeError, AttributeError):
+      logging.warning('Invalid JSON at line %d, skipped.', i)
+      continue
+
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=_COLUMN_NAMES,
+  )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard, output_fps)

community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/base.Dockerfile

@@ -0,0 +1,50 @@
+FROM python:3.9
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install basic libs.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        python3-opencv \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        vim \
+        screen \
+        libportaudio2 \
+        libusb-1.0-0-dev \
+        openjdk-17-jre
+
+# Add gcsfuse distribution URL as a package source and import its public key.
+RUN echo "deb https://packages.cloud.google.com/apt gcsfuse-`lsb_release -c -s` main" | sudo tee /etc/apt/sources.list.d/gcsfuse.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
+
+# Install gcsfuse.
+RUN apt-get update && apt-get install -y --no-install-recommends gcsfuse
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN ./google-cloud-sdk/install.sh -q
+# Make sure gsutil will use the default service account.
+RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
+
+
+# Install required libs.
+RUN pip install --upgrade pip
+RUN pip install pyyaml==5.4.1
+RUN pip install pycocotools==2.0.6
+RUN pip install opencv-python-headless==4.7.0.72
+RUN pip install numpy==1.24.2
+RUN pip install pandas==1.5.3
+RUN pip install Pillow==9.4.0
+RUN pip install apache-beam[gcp]==2.45.0
+RUN pip install object-detection==0.0.3
+RUN pip install google-cloud-storage==1.42.3
+RUN pip install gcsfs==2021.10.1
+RUN pip install pylint==2.17.2

community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/run.Dockerfile

@@ -0,0 +1,23 @@
+FROM gcr.io/automl-migration-test/automl-vision-data-converter-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+COPY model_oss/data_converter /automl_vision/data_converter
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision"
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["python3","data_converter/data_converter_main.py"]
+
+CMD ["--input_file_path=YOUR_INPUT_FILE",\
+"--input_file_type=csv",\
+"--objective=iod",\
+"--output_dir=YOUR_OUTPUT_DIR",\
+"--num_shard=10,10,10",\
+"--split_ratio=0.8,0.1,0.1"]

community-content/vertex_vision_model_garden/model_oss/diffusers/handler.py

@@ -4,9 +4,10 @@
 # pylint: disable=logging-fstring-interpolation
 
 import base64
+import io
 import logging
 import os
-from typing import Any, List, Tuple
+from typing import Any, List, Sequence, Tuple
 
 from diffusers import ControlNetModel
 from diffusers import DiffusionPipeline
@@ -20,6 +21,7 @@ from diffusers import StableDiffusionPipeline
 from diffusers import StableDiffusionUpscalePipeline
 from diffusers import TextToVideoZeroPipeline
 from diffusers import UniPCMultistepScheduler
+import imageio
 import numpy as np
 from PIL import Image
 import torch
@@ -43,6 +45,13 @@ TEXT_TO_VIDEO_ZERO_SHOT = "text-to-video-zero-shot"
 TEXT_TO_VIDEO = "text-to-video"
 
 
+def frames_to_video_bytes(frames: Sequence[np.ndarray], fps: int) -> bytes:
+  images = [Image.fromarray(array) for array in frames]
+  io_obj = io.BytesIO()
+  imageio.mimsave(io_obj, images, format=".mp4", fps=fps)
+  return io_obj.getvalue()
+
+
 class DiffusersHandler(BaseHandler):
   """Custom handler for TIMM models."""
 
@@ -214,7 +223,7 @@ class DiffusersHandler(BaseHandler):
         numpy_arrays = self.pipeline(prompt=prompt).images
         numpy_arrays = [(i * 255).astype("uint8") for i in numpy_arrays]
         videos.append(
-            video_format_converter.frames_to_video_bytes(numpy_arrays, fps=4)
+            frames_to_video_bytes(numpy_arrays, fps=4)
         )
       return videos
     elif self.task == TEXT_TO_VIDEO:
@@ -224,7 +233,7 @@ class DiffusersHandler(BaseHandler):
       # Therefore we need to split the output into different videos.
       predicted_images = np.array_split(predicted_images, len(prompts), axis=2)
       videos = [
-          video_format_converter.frames_to_video_bytes(images, fps=8)
+          frames_to_video_bytes(images, fps=8)
           for images in predicted_images
       ]
       return videos

community-content/vertex_vision_model_garden/model_oss/keras/dockerfile/serve.Dockerfile

@@ -0,0 +1,118 @@
+# Dockerfile for basic serving dockers with Keras.
+#
+# To build:
+# docker build -f model_oss/keras/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# Install basic libs.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        vim \
+        screen \
+        libtcmalloc-minimal4
+
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN ./google-cloud-sdk/install.sh -q
+# Make sure gsutil will use the default service account.
+RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
+
+
+# Install required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==5.4.1
+RUN pip install fsspec==2021.10.1
+RUN pip install gcsfs==2021.10.1
+RUN pip install tensorflow-text==2.11.0
+RUN pip install pyglove==0.1.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.2
+RUN pip install keras-cv==0.4.0
+RUN pip install tensorflow-datasets==4.8.3
+RUN pip install protobuf==3.20.3
+RUN pip install Pillow==9.5.0
+RUN pip install flask==2.3.2
+RUN pip install waitress==2.1.2
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+# Downloading gcloud package
+RUN curl https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz > /tmp/google-cloud-sdk.tar.gz
+
+# Installing the package
+RUN mkdir -p /usr/local/gcloud \
+  && tar -C /usr/local/gcloud -xvf /tmp/google-cloud-sdk.tar.gz \
+  && /usr/local/gcloud/google-cloud-sdk/install.sh
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Adding the package path to local
+ENV PATH $PATH:/usr/local/gcloud/google-cloud-sdk/bin
+
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8
+
+EXPOSE 8501
+
+WORKDIR /usr/local/lib/python3.8/dist-packages/official/vision
+
+COPY model_oss/keras /automl_vision/keras
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+ENV MODEL_PATH ""
+ENV IMAGE_WIDTH "512"
+ENV IMAGE_HEIGHT "512"
+
+COPY model_oss/keras/serve.py ./app.py
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["flask","run"]
+CMD ["--host=0.0.0.0", "--port=8501"]

community-content/vertex_vision_model_garden/model_oss/keras/dockerfile/train.Dockerfile

@@ -0,0 +1,111 @@
+# Dockerfile for basic training dockers with Keras.
+#
+# To build:
+# docker build -f model_oss/keras/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# Install basic libs.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        vim \
+        screen \
+        libtcmalloc-minimal4
+
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN ./google-cloud-sdk/install.sh -q
+# Make sure gsutil will use the default service account.
+RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
+
+
+# Install required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==5.4.1
+RUN pip install fsspec==2021.10.1
+RUN pip install gcsfs==2021.10.1
+RUN pip install tensorflow-text==2.11.0
+RUN pip install pyglove==0.1.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.2
+RUN pip install keras-cv==0.4.0
+RUN pip install tensorflow-datasets==4.8.3
+RUN pip install tensorflow-estimator==2.12.0
+RUN pip install tensorflow-gcs-config==2.12.0
+RUN pip install tensorflow-hub==0.13.0
+RUN pip install tensorflow-io-gcs-filesystem==0.32.0
+RUN pip install tensorflow-metadata==1.13.1
+RUN pip install tensorflow-probability==0.19.0
+RUN pip install tensorboard==2.12.2
+RUN pip install tensorboard-data-server==0.7.0
+RUN pip install tensorboard-plugin-wit==1.8.1
+RUN pip install protobuf==3.20.3
+RUN pip install pandas==1.5.3
+RUN pip install pandas-datareader==0.10.0
+RUN pip install pandas-gbq==0.17.9
+RUN pip install pycocotools==2.0.6
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8
+
+WORKDIR /usr/local/lib/python3.8/dist-packages/official/vision
+
+COPY model_oss/keras /automl_vision/keras
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+# Keras stable diffusion training codes set width and height as RESOLUTION.
+ENV RESOLUTION "512"
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["python3","keras/train.py"]

community-content/vertex_vision_model_garden/model_oss/keras/serve.py

@@ -0,0 +1,184 @@
+r"""Servers Keras Stable Diffusion models.
+
+python serve.py --model_path=<model path in gcs>
+
+curl -d \
+'{"prompt":"Hello Kitty"}' \
+-H "Content-Type: application/json" \
+-X POST http://localhost:8501/predict
+"""
+
+import base64
+import io
+import json
+import os
+from typing import List, Tuple
+
+from absl import app
+# The docker builds could not find flask and waitress.
+# pylint: disable=import-error
+from flask import Flask
+from flask import request
+from flask import Response
+import keras_cv
+from PIL import Image
+from waitress import serve
+
+from util import constants
+from util import fileutils
+
+
+flask_app = Flask(__name__)
+
+stable_diffusion_model = None
+
+
+model_path = os.environ.get('MODEL_PATH', '')
+if model_path.startswith(constants.GCS_URI_PREFIX):
+  print('Downloading models from gcs to local.')
+  os.makedirs(constants.LOCAL_MODEL_DIR, exist_ok=True)
+  fileutils.download_gcs_dir_to_local(
+      os.path.dirname(model_path), constants.LOCAL_MODEL_DIR
+  )
+  model_path = os.path.join(
+      constants.LOCAL_MODEL_DIR, os.path.basename(model_path)
+  )
+
+image_width = int(os.environ.get('IMAGE_WIDTH', 512))
+image_height = int(os.environ.get('IMAGE_HEIGHT', 512))
+
+print('image_width=', image_width, 'image_height=', image_height)
+print('Create Keras stable diffusion models.')
+stable_diffusion_model = keras_cv.models.StableDiffusion(
+    img_width=image_width,
+    img_height=image_height,
+    jit_compile=True,
+)
+
+if model_path:
+  # We just reload the weights of the fine-tuned diffusion model.
+  print('Initialize finetuned models from: ', model_path)
+  stable_diffusion_model.diffusion_model.load_weights(model_path)
+
+
+def error(message: str) -> str:
+  """Returns a JSON representing an error response."""
+  return json.dumps({
+      'success': False,
+      'error': message,
+  })
+
+
+def check_key_in_json(content: str, keys: List[str]) -> str:
+  for key in keys:
+    if key not in content:
+      return error('No {} in request {}.'.format(key, content))
+  return None
+
+
+def validate_json_key(json_key_string: str) -> Tuple[str, bool]:
+  try:
+    json_key = json.loads(json_key_string)
+  except (ValueError, TypeError):
+    return (error('Invalid key found in request'), False)
+  return (json_key, True)
+
+
+# The health check route is required for docker deployment in google cloud.
+@flask_app.route('/ping')
+def ping() -> Response:
+  """Health checks."""
+  return Response(status=200)
+
+
+# The return should be `Response` for docker deployment in google cloud.
+@flask_app.route('/predict', methods=['GET', 'POST'])
+def predict_model() -> Response:
+  """Predictions."""
+  if request.method == 'POST':
+    contents = request.get_json(force=True)
+
+    print('The input contents are:', contents)
+    batch_size = 1
+    num_steps = 25
+    seed = 1234
+    if 'parameters' in contents:
+      parameters = contents['parameters']
+      if 'batch_size' in parameters:
+        batch_size = int(parameters['batch_size'])
+      if 'num_steps' in parameters:
+        num_steps = int(parameters['num_steps'])
+      if 'seed' in parameters:
+        seed = int(parameters['seed'])
+    print('batch_size=', batch_size, 'num_steps=', num_steps, 'seed=', seed)
+    if batch_size < 1:
+      return Response(
+          response=error('The batch size must be a positive integar.'),
+          status=200,
+          mimetype='text/plain',
+      )
+    if num_steps < 1:
+      return Response(
+          response=error('The num steps must be a positive integar.'),
+          status=200,
+          mimetype='text/plain',
+      )
+    predictions = []
+    for content in contents['instances']:
+      print('Processing:', content)
+      prompt = content['prompt']
+      generated_image_array = stable_diffusion_model.text_to_image(
+          prompt=prompt,
+          batch_size=batch_size,
+          num_steps=num_steps,
+          seed=seed,
+      )
+
+      generated_image_bytes_array = []
+      for i in range(batch_size):
+        generated_image = Image.fromarray(generated_image_array[i])
+        # Converts the image to a base64-encoded string.
+        buffered_image = io.BytesIO()
+        generated_image.save(buffered_image, format='JPEG')
+        generated_image_bytes = base64.b64encode(
+            buffered_image.getvalue()
+        ).decode('utf-8')
+        generated_image_bytes_array.append(generated_image_bytes)
+      prediction = {
+          'prompt': prompt,
+          'predicted_image': generated_image_bytes_array,
+      }
+      predictions.append(prediction)
+
+    return Response(
+        response=json.dumps({
+            'success': True,
+            'predictions': predictions,
+        }),
+        status=200,
+        mimetype='text/plain',
+    )
+  else:
+    return Response(
+        response=json.dumps({
+            'success': True,
+            'isalive': stable_diffusion_model is not None,
+        }),
+        status=200,
+        mimetype='text/plain',
+    )
+
+
+def serve_main(unused_argv):
+  """The main function to serve Keras models."""
+  del unused_argv
+  # This is used when running locally only. When deploying to Google App
+  # Engine, a webserver process such as Gunicorn will serve the app.
+  # # Debug deployment.
+  # flask_app.run(host='0.0.0.0', port=8501, debug=True)
+  # Prod deployment.
+  serve(flask_app, host='0.0.0.0', port=8501)
+
+
+if __name__ == '__main__':
+  app.run(serve_main)

community-content/vertex_vision_model_garden/model_oss/keras/train.py

@@ -0,0 +1,363 @@
+"""Train Keras Stable Diffusion.
+
+Most the codes below are from
+https://keras.io/examples/generative/finetune_stable_diffusion/.
+"""
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+import keras_cv
+# pylint: disable=g-importing-member
+from keras_cv.models.stable_diffusion.clip_tokenizer import SimpleTokenizer
+from keras_cv.models.stable_diffusion.diffusion_model import DiffusionModel
+from keras_cv.models.stable_diffusion.image_encoder import ImageEncoder
+from keras_cv.models.stable_diffusion.noise_scheduler import NoiseScheduler
+from keras_cv.models.stable_diffusion.text_encoder import TextEncoder
+import numpy as np
+# The docker builds could not find pandas.
+# pylint: disable=import-error
+import pandas as pd
+import tensorflow as tf
+from tensorflow import keras
+import tensorflow.experimental.numpy as tnp
+
+from util import constants
+from util import fileutils
+
+_INPUT_CSV_PATH = flags.DEFINE_string(
+    'input_csv_path',
+    None,
+    'The input csv path.',
+    required=True,
+)
+
+_USE_MP = flags.DEFINE_bool(
+    'use_mp',
+    True,
+    'Enable mixed-precision training if the underlying GPU has tensor cores.',
+)
+
+_EPOCHS = flags.DEFINE_integer('epochs', 1, 'The number of epochs.')
+
+_OUTPUT_MODEL_DIR = flags.DEFINE_string(
+    'output_model_dir',
+    None,
+    'The output model dir.',
+    required=True,
+)
+
+# These hyperparameters defaults come from this tutorial by Hugging Face:
+# https://huggingface.co/docs/diffusers/training/text2image
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate', 1e-5, 'The learning rate parameter for AdamW optimizer.'
+)
+
+_BETA_1 = flags.DEFINE_float(
+    'beta_1', 0.9, 'The beta_1 parameter for AdamW optimizer.'
+)
+
+_BETA_2 = flags.DEFINE_float(
+    'beta_2', 0.999, 'The beta_2 parameter for AdamW optimizer.'
+)
+
+_WEIGHT_DECAY = flags.DEFINE_float(
+    'weight_decay', 1e-2, 'The weight decay parameter for AdamW optimizer.'
+)
+
+_EPSILON = flags.DEFINE_float(
+    'epsilon', 1e-08, 'The epsilon parameter for AdamW optimizer.'
+)
+
+RESOLUTION = int(os.environ.get('RESOLUTION', 512))
+
+# The padding token and maximum prompt length are specific to the text encoder.
+# If you're using a different text encoder be sure to change them accordingly.
+PADDING_TOKEN = 49407
+MAX_PROMPT_LENGTH = 77
+
+AUTO = tf.data.AUTOTUNE
+POS_IDS = tf.convert_to_tensor([list(range(MAX_PROMPT_LENGTH))], dtype=tf.int32)
+
+
+augmenter = keras.Sequential(
+    layers=[
+        keras_cv.layers.CenterCrop(RESOLUTION, RESOLUTION),
+        keras_cv.layers.RandomFlip(),
+        tf.keras.layers.Rescaling(scale=1.0 / 127.5, offset=-1),
+    ]
+)
+text_encoder = TextEncoder(MAX_PROMPT_LENGTH)
+
+
+def process_image(image_path, tokenized_text):
+  image = tf.io.read_file(image_path)
+  image = tf.io.decode_png(image, 3)
+  image = tf.image.resize(image, (RESOLUTION, RESOLUTION))
+  return image, tokenized_text
+
+
+def apply_augmentation(image_batch, token_batch):
+  return augmenter(image_batch), token_batch
+
+
+def run_text_encoder(image_batch, token_batch):
+  return (
+      image_batch,
+      token_batch,
+      text_encoder([token_batch, POS_IDS], training=False),
+  )
+
+
+def prepare_dict(image_batch, token_batch, encoded_text_batch):
+  return {
+      'images': image_batch,
+      'tokens': token_batch,
+      'encoded_text': encoded_text_batch,
+  }
+
+
+def prepare_dataset(image_paths, tokenized_texts, batch_size=1):
+  dataset = tf.data.Dataset.from_tensor_slices((image_paths, tokenized_texts))
+  dataset = dataset.shuffle(batch_size * 10)
+  dataset = dataset.map(process_image, num_parallel_calls=AUTO).batch(
+      batch_size
+  )
+  dataset = dataset.map(apply_augmentation, num_parallel_calls=AUTO)
+  dataset = dataset.map(run_text_encoder, num_parallel_calls=AUTO)
+  dataset = dataset.map(prepare_dict, num_parallel_calls=AUTO)
+  return dataset.prefetch(AUTO)
+
+
+def prepare_training_dataset(dataset_csv):
+  """Prepares training datasets."""
+  if dataset_csv.startswith(constants.GCS_URI_PREFIX):
+    if not os.path.exists(constants.LOCAL_DATA_DIR):
+      os.makedirs(constants.LOCAL_DATA_DIR)
+    logging.info(
+        'Start to download data from %s to %s.',
+        os.path.dirname(dataset_csv),
+        constants.LOCAL_DATA_DIR,
+    )
+    fileutils.download_gcs_dir_to_local(
+        os.path.dirname(dataset_csv), constants.LOCAL_DATA_DIR
+    )
+    data_frame = pd.read_csv(
+        os.path.join(constants.LOCAL_DATA_DIR, os.path.basename(dataset_csv))
+    )
+    data_frame['image_path'] = data_frame['image_path'].apply(
+        lambda x: os.path.join(constants.LOCAL_DATA_DIR, x)
+    )
+  else:
+    # Keeps the following codes for experiments with
+    # https://keras.io/examples/generative/finetune_stable_diffusion/.
+    data_path = tf.keras.utils.get_file(origin=dataset_csv, untar=True)
+    data_frame = pd.read_csv(os.path.join(data_path, 'data.csv'))
+    data_frame['image_path'] = data_frame['image_path'].apply(
+        lambda x: os.path.join(data_path, x)
+    )
+  data_frame.head()
+
+  # Load the tokenizer.
+  tokenizer = SimpleTokenizer()
+
+  #  Method to tokenize and pad the tokens.
+  def process_text(caption):
+    tokens = tokenizer.encode(caption)
+    tokens = tokens + [PADDING_TOKEN] * (MAX_PROMPT_LENGTH - len(tokens))
+    return np.array(tokens)
+
+  # Collate the tokenized captions into an array.
+  tokenized_texts = np.empty((len(data_frame), MAX_PROMPT_LENGTH))
+
+  all_captions = list(data_frame['caption'].values)
+  for i, caption in enumerate(all_captions):
+    tokenized_texts[i] = process_text(caption)
+
+  # Prepare the dataset.
+  training_dataset = prepare_dataset(
+      np.array(data_frame['image_path']), tokenized_texts, batch_size=4
+  )
+
+  return training_dataset
+
+
+class Trainer(tf.keras.Model):
+  """The trainer for Keras Stable Diffusion."""
+
+  # Reference:
+  # https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
+
+  def __init__(
+      self,
+      diffusion_model,
+      vae,
+      noise_scheduler,
+      use_mixed_precision=False,
+      max_grad_norm=1.0,
+      **kwargs,
+  ):
+    super().__init__(**kwargs)
+
+    self.diffusion_model = diffusion_model
+    self.vae = vae
+    self.noise_scheduler = noise_scheduler
+    self.max_grad_norm = max_grad_norm
+
+    self.use_mixed_precision = use_mixed_precision
+    self.vae.trainable = False
+
+  def train_step(self, inputs):
+    images = inputs['images']
+    encoded_text = inputs['encoded_text']
+    batch_size = tf.shape(images)[0]
+
+    with tf.GradientTape() as tape:
+      # Project image into the latent space and sample from it.
+      latents = self.sample_from_encoder_outputs(
+          self.vae(images, training=False)
+      )
+      # Know more about the magic number here:
+      # https://keras.io/examples/generative/fine_tune_via_textual_inversion/
+      latents = latents * 0.18215
+
+      # Sample noise that we'll add to the latents.
+      noise = tf.random.normal(tf.shape(latents))
+
+      # Sample a random timestep for each image.
+      timesteps = tnp.random.randint(
+          0, self.noise_scheduler.train_timesteps, (batch_size,)
+      )
+
+      # Add noise to the latents according to the noise magnitude at each
+      # timestep (this is the forward diffusion process).
+      noisy_latents = self.noise_scheduler.add_noise(
+          tf.cast(latents, noise.dtype), noise, timesteps
+      )
+
+      # Get the target for loss depending on the prediction type
+      # just the sampled noise for now.
+      target = noise  # noise_schedule.predict_epsilon == True
+
+      # Predict the noise residual and compute loss.
+      # pylint: disable=unnecessary-lambda
+      timestep_embedding = tf.map_fn(
+          lambda t: self.get_timestep_embedding(t), timesteps, dtype=tf.float32
+      )
+      timestep_embedding = tf.squeeze(timestep_embedding, 1)
+      model_pred = self.diffusion_model(
+          [noisy_latents, timestep_embedding, encoded_text], training=True
+      )
+      loss = self.compiled_loss(target, model_pred)
+      if self.use_mixed_precision:
+        loss = self.optimizer.get_scaled_loss(loss)
+
+    # Update parameters of the diffusion model.
+    trainable_vars = self.diffusion_model.trainable_variables
+    gradients = tape.gradient(loss, trainable_vars)
+    if self.use_mixed_precision:
+      gradients = self.optimizer.get_unscaled_gradients(gradients)
+    gradients = [tf.clip_by_norm(g, self.max_grad_norm) for g in gradients]
+    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
+
+    return {m.name: m.result() for m in self.metrics}
+
+  def get_timestep_embedding(self, timestep, dim=320, max_period=10000):
+    half = dim // 2
+    log_max_preiod = tf.math.log(tf.cast(max_period, tf.float32))
+    # The docker builds could not support unary `-`.
+    # pylint: disable=invalid-unary-operand-type
+    freqs = tf.math.exp(
+        -log_max_preiod * tf.range(0, half, dtype=tf.float32) / half
+    )
+    args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
+    embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
+    embedding = tf.reshape(embedding, [1, -1])
+    return embedding
+
+  def sample_from_encoder_outputs(self, outputs):
+    mean, logvar = tf.split(outputs, 2, axis=-1)
+    logvar = tf.clip_by_value(logvar, -30.0, 20.0)
+    std = tf.exp(0.5 * logvar)
+    sample = tf.random.normal(tf.shape(mean), dtype=mean.dtype)
+    return mean + std * sample
+
+  def save_weights(
+      self, filepath, overwrite=True, save_format=None, options=None
+  ):
+    # Overriding this method will allow us to use the `ModelCheckpoint`
+    # callback directly with this trainer class. In this case, it will
+    # only checkpoint the `diffusion_model` since that's what we're training
+    # during fine-tuning.
+    self.diffusion_model.save_weights(
+        filepath=filepath,
+        overwrite=overwrite,
+        save_format=save_format,
+        options=options,
+    )
+
+
+def main(_) -> None:
+  # _INPUT_CSV_PATH and _OUTPUT_MODEL_DIR should have the format as
+  # gs://<bucket_name>/<object_name>.
+  if _INPUT_CSV_PATH.value:
+    if not _INPUT_CSV_PATH.value.startswith(constants.GCS_URI_PREFIX):
+      raise ValueError('The input csv path should be a gcs path like gs://<>')
+  if _OUTPUT_MODEL_DIR.value:
+    if not _OUTPUT_MODEL_DIR.value.startswith(constants.GCS_URI_PREFIX):
+      raise ValueError('The output model dir should be a gcs path like gs://<>')
+
+  if _USE_MP.value:
+    keras.mixed_precision.set_global_policy('mixed_float16')
+
+  image_encoder = ImageEncoder(RESOLUTION, RESOLUTION)
+  diffusion_ft_trainer = Trainer(
+      diffusion_model=DiffusionModel(RESOLUTION, RESOLUTION, MAX_PROMPT_LENGTH),
+      # Remove the top layer from the encoder, which cuts off the variance and
+      # only returns the mean.
+      vae=tf.keras.Model(
+          image_encoder.input,
+          image_encoder.layers[-2].output,
+      ),
+      noise_scheduler=NoiseScheduler(),
+      use_mixed_precision=_USE_MP.value,
+  )
+
+  optimizer = tf.keras.optimizers.experimental.AdamW(
+      learning_rate=_LEARNING_RATE.value,
+      weight_decay=_WEIGHT_DECAY.value,
+      beta_1=_BETA_1.value,
+      beta_2=_BETA_2.value,
+      epsilon=_EPSILON.value,
+  )
+  diffusion_ft_trainer.compile(optimizer=optimizer, loss='mse')
+
+  training_dataset = prepare_training_dataset(_INPUT_CSV_PATH.value)
+
+  # Note: gcsfuse does not work for Keras. We saves the trained models locally
+  # first, and then copy to gcs storages.
+  if not os.path.exists(constants.LOCAL_MODEL_DIR):
+    os.makedirs(constants.LOCAL_MODEL_DIR)
+  # The default saved model is in HDF5.
+  ckpt_path = os.path.join(constants.LOCAL_MODEL_DIR, 'saved_model.h5')
+  ckpt_callback = tf.keras.callbacks.ModelCheckpoint(
+      ckpt_path,
+      save_weights_only=True,
+      monitor='loss',
+      mode='min',
+  )
+  diffusion_ft_trainer.fit(
+      training_dataset, epochs=_EPOCHS.value, callbacks=[ckpt_callback]
+  )
+
+  # Copies the files in constants.LOCAL_MODEL_DIR to output_model_dir.
+  fileutils.upload_local_dir_to_gcs(
+      constants.LOCAL_MODEL_DIR, _OUTPUT_MODEL_DIR.value
+  )
+
+  return
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/lm-evaluation-harness/dockerfile/eval.Dockerfile

@@ -0,0 +1,40 @@
+# Dockerfile for lm-evaluation-harness evaluation.
+#
+# To build:
+# docker build -f model_oss/lm-evaluation-harness/dockerfile/eval.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/{YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/{YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
+
+USER root
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+RUN apt-get install -y --no-install-recommends apt-utils
+RUN apt-get install -y --no-install-recommends curl
+RUN apt-get install -y --no-install-recommends wget
+RUN apt-get install -y --no-install-recommends git
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
+
+# Install lm-evaluation-harness
+RUN git clone https://github.com/EleutherAI/lm-evaluation-harness
+WORKDIR lm-evaluation-harness
+# Pin version up to date 08/08/2023
+RUN git reset --hard b952a206de210b72b1bf750fbab38c26121e0dc0
+# Edit tokenizer loading function to avoid using fast tokenizer for OpenLLaMA
+RUN sed -i '355 i\        use_fast = not pretrained.startswith("openlm-research/open_llama")' lm_eval/models/huggingface.py
+RUN sed -i '360 i\            use_fast=use_fast,' lm_eval/models/huggingface.py
+# Install from source while including the sentencepiece dependency
+RUN pip install -e ".[sentencepiece]"

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/base.Dockerfile

@@ -0,0 +1,64 @@
+FROM tensorflow/build:2.12-python3.9
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# Install basic libs.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        vim \
+        libtcmalloc-minimal4
+
+
+# Install google cloud CLI.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-430.0.0-linux-x86.tar.gz
+RUN tar xzf google-cloud-cli-430.0.0-linux-x86.tar.gz
+RUN ./google-cloud-sdk/install.sh -q
+# Make sure gsutil will use the default service account.
+RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
+
+
+# Install required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==6.0
+RUN pip install fsspec==2023.4.0
+RUN pip install gcsfs==2023.4.0
+RUN pip install tf-models-official==2.12.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.3
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/model_export.Dockerfile

@@ -0,0 +1,13 @@
+FROM gcr.io/automl-migration-test/movinet-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN wget https://raw.githubusercontent.com/tensorflow/models/954dd73bffd43174bd3ca26a4a34abebe4147570/official/projects/movinet/tools/export_saved_model.py \
+  -O /usr/local/lib/python3.9/dist-packages/official/projects/movinet/tools/export_saved_model.py
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+ENTRYPOINT ["python3", "-m", "official.projects.movinet.tools.export_saved_model"]

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/serve.Dockerfile

@@ -0,0 +1,18 @@
+FROM gcr.io/automl-migration-test/movinet-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN pip install flask==2.3.2
+RUN pip install waitress==2.1.2
+
+RUN mkdir -p /automl_vision/movinet/serving
+COPY model_oss/movinet/serving /automl_vision/movinet/serving
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+ENTRYPOINT ["flask", "--app", "movinet.serving.serving_main", "run"]
+CMD ["--host=0.0.0.0", "--port=8501"]

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/train.Dockerfile

@@ -0,0 +1,18 @@
+FROM gcr.io/automl-migration-test/movinet-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN mkdir -p /automl_vision/movinet
+COPY model_oss/movinet/*.py /automl_vision/movinet/
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["python3","movinet/train.py"]

community-content/vertex_vision_model_garden/model_oss/movinet/serving/serving_main.py

@@ -0,0 +1,142 @@
+"""Main executable for MoViNet online / batch predictions."""
+
+from collections.abc import Sequence
+import json
+import os
+
+from absl import app
+from absl import logging
+import flask
+import tensorflow as tf
+import waitress
+
+from movinet.serving import video_serving_lib
+from util import constants
+
+
+flask_app = flask.Flask(__name__)
+logging.set_verbosity(logging.INFO)
+
+movinet_model = None
+
+_BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '1'))
+_NUM_FRAMES = int(os.environ.get('NUM_FRAMES', '32'))
+_FPS = float(os.environ.get('FPS', '5'))
+_OVERLAP_FRAMES = int(os.environ.get('OVERLAP_FRAMES', '24'))
+_OBJECTIVE = os.environ.get(
+    'OBJECTIVE', constants.OBJECTIVE_VIDEO_CLASSIFICATION
+).lower()
+
+# VAR parameters.
+_CONFIDENCE_THRESHOLD = float(os.environ.get('CONFIDENCE_THRESHOLD', '0.5'))
+_MIN_GAP_TIME = float(os.environ.get('MIN_GAP_TIME', '1.5'))
+
+
+def load_movinet_model() -> None:
+  model_path = os.environ.get('MODEL_PATH')
+
+  if not model_path:
+    raise app.UsageError('Missing MODEL_PATH environment variable.')
+
+  # We just reload the weights of the fine-tuned diffusion model.
+  logging.info('Initialize finetuned models from: %s', model_path)
+  global movinet_model
+  movinet_model = tf.saved_model.load(model_path)
+
+
+load_movinet_model()
+
+
+def error(message: str) -> str:
+  """Returns a JSON representing an error response."""
+  return json.dumps({
+      'success': False,
+      'error': message,
+  })
+
+
+# The health check route is required for docker deployment in google cloud.
+@flask_app.route('/ping')
+def ping() -> flask.Response:
+  """Health checks."""
+  return flask.Response(status=200)
+
+
+# The return should be `Response` for docker deployment in google cloud.
+@flask_app.route('/predict', methods=['GET', 'POST'])
+def predict_model() -> flask.Response:
+  """Predictions."""
+  if flask.request.method == 'POST':
+    contents = flask.request.get_json(force=True)
+
+    logging.info('The input contents are: %s', contents)
+    instances = contents.get('instances', [])
+
+    try:
+      predictions = []
+      for instance in instances:
+        executor = video_serving_lib.parse_request(instance)
+        prediction = executor.get_prediction(
+            movinet_model,
+            _BATCH_SIZE,
+            _FPS,
+            _NUM_FRAMES,
+            _OVERLAP_FRAMES,
+            _OBJECTIVE,
+        )
+        if _OBJECTIVE == constants.OBJECTIVE_VIDEO_CLASSIFICATION:
+          prediction = video_serving_lib.postprocess_vcn(prediction)
+        elif _OBJECTIVE == constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION:
+          prediction = video_serving_lib.postprocess_var(
+              executor.windows, prediction, _CONFIDENCE_THRESHOLD, _MIN_GAP_TIME
+          )
+        predictions.append(prediction)
+    except ValueError as e:
+      return flask.Response(
+          error(str(e)), status=500, mimetype='application/json'
+      )
+
+    return flask.Response(
+        response=json.dumps({
+            'success': True,
+            'predictions': predictions,
+        }),
+        status=200,
+        mimetype='application/json',
+    )
+  else:
+    return flask.Response(
+        response=json.dumps({
+            'success': True,
+            'isalive': movinet_model is not None,
+        }),
+        status=200,
+        mimetype='application/json',
+    )
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+  # This is used when running locally only. When deploying to Google App
+  # Engine, a webserver process such as Gunicorn will serve the app.
+  # # Debug deployment.
+  # flask_app.run(host='0.0.0.0', port=8501, debug=True)
+  # Prod deployment.
+  if _OBJECTIVE not in [
+      constants.OBJECTIVE_VIDEO_CLASSIFICATION,
+      constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION,
+  ]:
+    raise app.UsageError('Objective must be vcn or var.')
+  logging.info(
+      'Env: batch_size: %s, num_frames: %s, fps: %s, overlap_frames: %s',
+      _BATCH_SIZE,
+      _NUM_FRAMES,
+      _FPS,
+      _OVERLAP_FRAMES,
+  )
+  waitress.serve(flask_app, host='0.0.0.0', port=8501)
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/movinet/serving/video_serving_lib.py

@@ -0,0 +1,462 @@
+"""Lib for handling video prediction requests.
+
+The VCN inference algorithm is as follows:
+1. Find all video frames within the given clip according to the sampling FPS.
+2. Create possibly overlapping sliding windows according to the num_frames and
+overlap_frames parameters. The last window might have a larger overlap if it
+doesn't exactly fit.
+3. Run model inference on each sliding window and compute softmax to obtain
+probabilities.
+4. Average the probabilities over all sliding windows.
+
+The VAR inference algorithm is very similar to VCN, with a few differences:
+1. The last sliding window is discarded if it does not exactly fit.
+2. Instead of averaging, the postprocessing consists of temporal nonmaximal
+suppression and removing background and low-confidence labels.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import os
+from typing import Any, Dict, Optional, Sequence, Union, cast
+
+from absl import logging
+import cv2
+import numpy as np
+import tensorflow as tf
+
+from util import constants
+from util import fileutils
+
+
+_JSON_LABEL_KEY = 'label'
+_JSON_GCS_URI_KEY = 'content'
+_JSON_CONFIDENCE_KEY = 'confidence'
+_JSON_START_TIME_KEY = 'timeSegmentStart'
+_JSON_END_TIME_KEY = 'timeSegmentEnd'
+_BACKGROUND_LABEL = 0
+_JSON_REQUIRED_KEYS = [
+    _JSON_GCS_URI_KEY,
+    _JSON_START_TIME_KEY,
+    _JSON_END_TIME_KEY,
+]
+_IMAGE_WIDTH = int(os.environ.get('IMAGE_WIDTH', '172'))
+_IMAGE_HEIGHT = int(os.environ.get('IMAGE_HEIGHT', '172'))
+
+
+@dataclasses.dataclass
+class DetectionOutput:
+  timestamp: float
+  label: int
+  confidence: float
+
+  def to_json_obj(self) -> Dict[str, Union[int, float]]:
+    """Encodes self as a dict for JSON serialization."""
+    return {
+        _JSON_LABEL_KEY: self.label,
+        _JSON_START_TIME_KEY: self.timestamp,
+        _JSON_END_TIME_KEY: self.timestamp,
+        _JSON_CONFIDENCE_KEY: self.confidence,
+    }
+
+
+def create_detection_output(
+    timestamp: float, predictions: np.ndarray
+) -> DetectionOutput:
+  label = np.argmax(predictions).item()
+  confidence: float = predictions[label].item()
+  return DetectionOutput(timestamp, label, confidence)
+
+
+class SlidingWindow:
+  """Represents a sliding window with start / end timestamps."""
+
+  def __init__(self, fps: float, frames: Sequence[int]):
+    if not frames:
+      raise ValueError('Sliding window cannot be empty.')
+    self.frames = frames
+    self.start_time = frames[0] / fps
+    self.end_time = frames[-1] / fps
+    self.frame_data: list[Optional[np.ndarray]] = []
+    self.clear_frame_data()
+
+  def load_cache_from(self, other: SlidingWindow) -> int:
+    """Loads cache from another sliding window if possible."""
+    cache_count = 0
+    for i, frame in enumerate(self.frames):
+      try:
+        other_idx = other.frames.index(frame)
+        self.frame_data[i] = other.frame_data[other_idx]
+        cache_count += 1
+      except ValueError:
+        # Cache miss.
+        pass
+    return cache_count
+
+  def load_frames(self, video: Any) -> Sequence[np.ndarray]:
+    """Loads frames of this sliding window from a video."""
+    for i, frame in enumerate(self.frames):
+      if self.frame_data[i] is None:
+        video.set(cv2.CAP_PROP_POS_FRAMES, frame)
+        ret, frame = video.read()
+        if not ret:
+          raise IOError(f'Failed to read video at frame {frame}.')
+        self.frame_data[i] = cv2.resize(frame, (_IMAGE_WIDTH, _IMAGE_HEIGHT))
+    return cast(Sequence[np.ndarray], self.frame_data)
+
+  def clear_frame_data(self) -> None:
+    """Clears frame data of this sliding window to reduce memory usage."""
+    self.frame_data: list[Optional[np.ndarray]] = [None] * len(self)
+
+  def __len__(self) -> int:
+    return len(self.frames)
+
+  @property
+  def middle_timestamp(self) -> float:
+    return (self.start_time + self.end_time) / 2
+
+
+def _get_sliding_windows(
+    frames: Sequence[int],
+    original_fps: float,
+    window_size: int,
+    overlap: int,
+    flush_last_window: bool,
+) -> Sequence[SlidingWindow]:
+  """Computes a list of sliding windows from frames.
+
+  Args:
+    frames: A list of frame indices.
+    original_fps: Frames per second of the original video.
+    window_size: Number of frames in a single window.
+    overlap: Number of overlapping frames in adjacent windows.
+    flush_last_window: Where to flush the last window if there are not enough
+      frames left.
+
+  Returns:
+    A list of sliding windows, each has a list of frame indices. The last two
+    windows might have a larger overlap if the last window does not exactly fit
+    and flush_last_window is set to True.
+
+  Raises:
+    ValueError: Arguments are invalid.
+  """
+  if window_size <= overlap:
+    raise ValueError(f'Window size {window_size} <= overlap {overlap}')
+  total_frames = len(frames)
+  windows: list[SlidingWindow] = []
+  for i in range(0, total_frames, window_size - overlap):
+    if i == 0 or i + window_size <= total_frames:
+      windows.append(SlidingWindow(original_fps, frames[i : i + window_size]))
+    elif i + overlap < total_frames and flush_last_window:
+      # Some frames in this window are not covered by the previous window.
+      windows.append(
+          SlidingWindow(
+              original_fps, frames[total_frames - window_size : total_frames]
+          )
+      )
+  return windows
+
+
+def _sample_frame_indices(
+    start_time: float,
+    end_time: float,
+    original_fps: float,
+    sample_fps: float,
+    max_frames: int,
+    padding_left: int = 0,
+    padding_right: int = 0,
+) -> Sequence[int]:
+  """Samples frames from start_time to end_time by sample_fps.
+
+  Args:
+    start_time: Start timestamp in seconds.
+    end_time: End timestamp in seconds.
+    original_fps: Frames per second of the original video.
+    sample_fps: Number of frames to sample per second.
+    max_frames: Total number of frames in the video.
+    padding_left: Padding to add to the start in frames. Padded frames will be
+      duplicates of the first frame.
+    padding_right: Padding to add to the end in frames. Padded frames will be
+      duplicates of the last frame.
+
+  Returns:
+    A list of sampled frame indices.
+  """
+  ret = [
+      min(max_frames - 1, round(t * original_fps))
+      for t in np.arange(start_time, end_time, 1 / sample_fps)
+  ]
+  if ret:
+    ret = [ret[0]] * padding_left + ret + [ret[-1]] * padding_right
+  return ret
+
+
+class VideoPredictionExecutor:
+  """Represents a Video prediction request with a video clip."""
+
+  def __init__(self, gcs_uri: str, start_time: float, end_time: float):
+    self._gcs_uri = gcs_uri
+    self._start_time = start_time
+    self._end_time = end_time
+    self.windows: Sequence[SlidingWindow] = []
+    self._last_window: SlidingWindow = None
+
+  def _read_frames_from_window(
+      self, video: Any, new_window: SlidingWindow
+  ) -> Sequence[np.ndarray]:
+    """Reads video frames from the new window.
+
+    Args:
+      video: Video loaded with cv2.
+      new_window: A list of sorted frame indices in the new window.
+
+    Returns:
+      Frame data from the video as a list of numpy arrays.
+
+    Raises:
+      IOError: Failed to read video.
+    """
+    # Caches frames as much as possible.
+    if self._last_window is not None:
+      cache_count = new_window.load_cache_from(self._last_window)
+      logging.info('Cached %d frames.', cache_count)
+      self._last_window.clear_frame_data()
+    self._last_window = new_window
+    return new_window.load_frames(video)
+
+  def _predict(
+      self, model: Any, video: Any, batched_windows: Sequence[SlidingWindow]
+  ) -> np.ndarray:
+    """Run model inference on specific frames of a video.
+
+    Args:
+      model: MoViNet model.
+      video: Video loaded with cv2.
+      batched_windows: A batch of sliding windows to predict. Each element is an
+        integer frame index. Must have equal number of frames in each window.
+
+    Returns:
+      Prediction results.
+
+    Raises:
+      ValueError: Batched windows are not sorted, or do not have equal number of
+        frames in each window.
+      IOError: Failed to read video.
+    """
+    if any(
+        (
+            len(window) != len(batched_windows[0])
+            for window in batched_windows[1:]
+        )
+    ):
+      raise ValueError(
+          'Batched windows do not have equal number of frames in each window.'
+      )
+    batch = []
+    logging.info('Loading video frames...')
+    for window in batched_windows:
+      logging.info('Predict frames: %s', window.frames)
+      frames = self._read_frames_from_window(video, window)
+      batch.append(frames)
+    input_tensor = tf.convert_to_tensor(batch, dtype=tf.float32) / 255.0
+    logging.info('Predict: Input tensor shape %s', input_tensor.shape)
+    predictions = model({'image': input_tensor})
+    logging.info('Running softmax on predictions...')
+    predictions = tf.nn.softmax(predictions, axis=1)
+    return predictions.numpy()
+
+  def get_prediction(
+      self,
+      model: Any,
+      batch_size: int,
+      fps: float,
+      num_frames: int,
+      overlap_frames: int,
+      objective: str,
+  ) -> Sequence[np.ndarray]:
+    """Predicts the video clip with the model.
+
+    Args:
+      model: The loaded MoViNet model.
+      batch_size: Batch size for prediction.
+      fps: Video sampling FPS.
+      num_frames: Number of frames in a single predictions. If the model is
+        exported with a fixed input shape, this must match its num_frames
+        dimension.
+      overlap_frames: Number of overlapping frames of consecutive sliding
+        windows.
+      objective: A string `vcn` or `var`.
+
+    Returns:
+      A list of floats as the prediction response.
+
+    Raises:
+      IOError: The video fails to load.
+      ValueError: Some arguments are invalid.
+    """
+    if objective not in [
+        constants.OBJECTIVE_VIDEO_CLASSIFICATION,
+        constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION,
+    ]:
+      raise ValueError(f'{objective} objective is not supported.')
+
+    # cv2 expects a local path so we need to download the video from GCS.
+    local_file_path = fileutils.generate_tmp_path(
+        os.path.splitext(self._gcs_uri)[1]
+    )
+    logging.info('Downloading %s to %s...', self._gcs_uri, local_file_path)
+    fileutils.download_gcs_file_to_local(self._gcs_uri, local_file_path)
+    logging.info('Download %s complete.', self._gcs_uri)
+
+    # Loads video.
+    video = cv2.VideoCapture(local_file_path)
+    total_frames = video.get(cv2.CAP_PROP_FRAME_COUNT)
+    original_fps = video.get(cv2.CAP_PROP_FPS)
+    if not original_fps:
+      # 0 or None indicates the video is invalid.
+      raise IOError(f'Failed to load {self._gcs_uri}.')
+    video_length = total_frames / original_fps
+    self._start_time = max(0, self._start_time)
+    self._end_time = min(video_length, self._end_time)
+    padding = (
+        (num_frames // 2)
+        if objective == constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION
+        else 0
+    )
+
+    # Computes sliding windows.
+    frame_indices = _sample_frame_indices(
+        self._start_time,
+        self._end_time,
+        original_fps,
+        fps,
+        total_frames,
+        padding,
+        padding,
+    )
+    logging.info('Frame indices: %s', frame_indices)
+    self.windows = _get_sliding_windows(
+        frame_indices,
+        original_fps,
+        num_frames,
+        overlap_frames,
+        objective != 'var',
+    )
+    if not self.windows:
+      raise ValueError(
+          f'No sliding windows found from {self._start_time} to'
+          f' {self._end_time}.'
+      )
+    self._last_window = None
+
+    # Runs inference.
+    predictions = []
+    for i in range(0, len(self.windows), batch_size):
+      predictions.extend(
+          self._predict(model, video, self.windows[i : i + batch_size])
+      )
+    return predictions
+
+
+def parse_request(req_json: Any) -> VideoPredictionExecutor:
+  """Parses VideoPredictionExecutor from request JSON object.
+
+  Args:
+    req_json: Request JSON object.
+
+  Returns:
+    Parsed VideoPredictionExecutor.
+
+  Raises:
+    ValueError: Request JSON object is invalid.
+  """
+  for key in _JSON_REQUIRED_KEYS:
+    if key not in req_json:
+      raise ValueError(f'{key} not found in {req_json}.')
+  gcs_uri = req_json[_JSON_GCS_URI_KEY]
+  start_time = float(req_json[_JSON_START_TIME_KEY].removesuffix('s'))
+  end_time = float(req_json[_JSON_END_TIME_KEY].removesuffix('s'))
+  return VideoPredictionExecutor(gcs_uri, start_time, end_time)
+
+
+def postprocess_vcn(predictions: Sequence[np.ndarray]) -> Sequence[float]:
+  """Aggregates VCN predictions of sliding windows."""
+  return np.mean(predictions, axis=0).tolist()
+
+
+def temporal_nonmaximal_suppression(
+    detections: Sequence[DetectionOutput], min_gap_time: float
+) -> Sequence[DetectionOutput]:
+  """Nonmaximal suppression for key frame detection.
+
+  For consecutive packets of the same label within a pre-defined duration, we
+  only keep the one with the highest confidence score. Such duration can be
+  determined by performing data analysis on users' dataset.
+
+  Args:
+    detections: A list of DetectionOutputs.
+    min_gap_time: Minimum time between consecutive key frames of the same label
+      in seconds.
+
+  Returns:
+    DetectionOutput after nonmaximal suppression sorted in ascending timestamps.
+  """
+  max_label = max([detection.label for detection in detections])
+  prev_detections: list[Optional[DetectionOutput]] = [None] * (max_label + 1)
+  ret: list[DetectionOutput] = []
+  by_time = lambda x: x.timestamp
+  for detection in sorted(detections, key=by_time):
+    prev_detection = prev_detections[detection.label]
+    prev_detections[detection.label] = detection
+    if not prev_detection:
+      continue
+    if detection.timestamp - prev_detection.timestamp > min_gap_time:
+      ret.append(prev_detection)
+      continue
+    detection.confidence = max(detection.confidence, prev_detection.confidence)
+  ret.extend((d for d in prev_detections if d is not None))
+  return sorted(ret, key=by_time)
+
+
+def postprocess_var(
+    windows: Sequence[SlidingWindow],
+    predictions: Sequence[np.ndarray],
+    confidence_threshold: float,
+    min_gap_time: float,
+) -> Sequence[Dict[str, Any]]:
+  """Generates a list of detected keyframes from sliding window predictions.
+
+  Args:
+    windows: Sliding windows.
+    predictions: A list of predictions of sliding windows.
+    confidence_threshold: Only probabilities greater than this threshold will
+      contribute to the final result.
+    min_gap_time: Minimum time between consecutive key frames of the same label
+      in seconds. Used in temporal nonmaximal suppression.
+
+  Returns:
+    A sequence of dictionaries, each item has the following keys:
+    - label: Integer label of the detection result.
+    - timeSegmentStart: Start timestamp in seconds.
+    - timeSegmentEnd: End timestamp in seconds. Always equals timeSegmentStart.
+  """
+  if len(windows) != len(predictions):
+    raise ValueError('Mismatched # of windows with # of predictions.')
+
+  # Creates detection results from windows, filtering out the background label.
+  detections = [
+      create_detection_output(window.middle_timestamp, predictions[i])
+      for i, window in enumerate(windows)
+  ]
+
+  # Temporal nonmaximal suppression.
+  detections = temporal_nonmaximal_suppression(detections, min_gap_time)
+
+  # Filters out ones with low confidence and the background label.
+  return [
+      x.to_json_obj()
+      for x in detections
+      if x.label != _BACKGROUND_LABEL and x.confidence > confidence_threshold
+  ]

community-content/vertex_vision_model_garden/model_oss/movinet/train.py

@@ -0,0 +1,210 @@
+"""Main executable for MoViNet docker."""
+
+import json
+import os
+from typing import Sequence, Any
+
+from absl import app
+from absl import flags
+from absl import logging
+import gin
+import hypertune
+import tensorflow as tf
+
+from util import constants
+from util import hypertune_utils
+from official.common import distribute_utils
+from official.common import flags as tfm_flags
+from official.core import task_factory
+from official.core import train_lib
+from official.core import train_utils
+from official.modeling import performance
+# Import movinet libraries to register the backbone and model into tf.vision
+# model garden factory.
+# pylint: disable=unused-import
+from official.projects.movinet.modeling import movinet
+from official.projects.movinet.modeling import movinet_model
+from official.vision import registry_imports
+# pylint: enable=unused-import
+
+
+FLAGS = flags.FLAGS
+
+_FILE_TYPE_TFRECORD = 'tfrecord'
+
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate', None, 'The learning rate of this training job.'
+)
+
+_NUM_CLASSES = flags.DEFINE_integer(
+    'num_classes', None, 'The number of classes.'
+)
+
+_INIT_CHECKPOINT = flags.DEFINE_string(
+    'init_checkpoint', None, 'The initial checkpoint of this training job.'
+)
+
+_INPUT_TRAIN_DATA_PATH = flags.DEFINE_string(
+    'input_train_data_path', None, 'Input train data path.'
+)
+
+_INPUT_VALIDATION_DATA_PATH = flags.DEFINE_string(
+    'input_validation_data_path', None, 'Input validation data path.'
+)
+
+_GLOBAL_BATCH_SIZE = flags.DEFINE_integer(
+    'global_batch_size', None, 'Global batch size.'
+)
+
+_PREFETCH_BUFFER_SIZE = flags.DEFINE_integer(
+    'prefetch_buffer_size', None, 'Prefetch buffer size.'
+)
+
+_SHUFFLE_BUFFER_SIZE = flags.DEFINE_integer(
+    'shuffle_buffer_size', None, 'Shuffle buffer size.'
+)
+
+_TRAIN_STEPS = flags.DEFINE_integer('train_steps', None, 'Train steps.')
+_LOG_LEVEL = flags.DEFINE_enum(
+    'log_level',
+    'INFO',
+    ['FATAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG'],
+    'Log level.',
+)
+
+
+def parse_params() -> Any:
+  """Parses parameters."""
+  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params)
+  params = train_utils.parse_configuration(FLAGS, lock_return=False)
+  if _INIT_CHECKPOINT.value:
+    params.task.init_checkpoint = _INIT_CHECKPOINT.value
+    params.task.init_checkpoint_modules = 'backbone'
+  if _NUM_CLASSES.value:
+    params.task.model.num_classes = _NUM_CLASSES.value
+    params.task.train_data.num_classes = _NUM_CLASSES.value
+    params.task.validation_data.num_classes = _NUM_CLASSES.value
+  # If users set input train/validation data path, we assume the data are
+  # converted from data converter as tfrecord. Users can use tfds by writing
+  # their own config directly, and no need to override this parameter.
+  if _INPUT_TRAIN_DATA_PATH.value:
+    params.task.train_data.input_path = _INPUT_TRAIN_DATA_PATH.value
+    params.task.train_data.file_type = _FILE_TYPE_TFRECORD
+    params.task.train_data.tfds_name = ''
+  if _INPUT_VALIDATION_DATA_PATH.value:
+    params.task.validation_data.input_path = _INPUT_VALIDATION_DATA_PATH.value
+    params.task.validation_data.file_type = _FILE_TYPE_TFRECORD
+    params.task.validation_data.tfds_name = ''
+  if _GLOBAL_BATCH_SIZE.value:
+    params.task.train_data.global_batch_size = _GLOBAL_BATCH_SIZE.value
+    params.task.validation_data.global_batch_size = _GLOBAL_BATCH_SIZE.value
+  if _PREFETCH_BUFFER_SIZE.value:
+    params.task.train_data.prefetch_buffer_size = _PREFETCH_BUFFER_SIZE.value
+    params.task.validation_data.prefetch_buffer_size = (
+        _PREFETCH_BUFFER_SIZE.value
+    )
+  if _SHUFFLE_BUFFER_SIZE.value:
+    params.task.train_data.shuffle_buffer_size = _SHUFFLE_BUFFER_SIZE.value
+  if _TRAIN_STEPS.value:
+    params.trainer.train_steps = _TRAIN_STEPS.value
+  if _LEARNING_RATE.value:
+    logging.info('Updating learning_rate: %s', _LEARNING_RATE.value)
+    # Use `get` method of train_utils.hyperparams.OneOfConfig to get learning
+    # rate config.
+    learning_rate = params.trainer.optimizer_config.learning_rate.get()
+    if hasattr(learning_rate, 'initial_learning_rate'):
+      learning_rate.initial_learning_rate = _LEARNING_RATE.value
+    else:
+      logging.warning('Cannot set learning rate for %s', learning_rate)
+  # Set default params for best checkpoints.
+  params.trainer.best_checkpoint_export_subdir = constants.BEST_CKPT_DIRNAME
+  params.trainer.best_checkpoint_metric_comp = constants.BEST_CKPT_METRIC_COMP
+  params.trainer.best_checkpoint_eval_metric = (
+      constants.VIDEO_CLASSIFICATION_BEST_EVAL_METRIC
+  )
+  return params
+
+
+def main(argv: Sequence[str]) -> None:
+  logging.set_verbosity(_LOG_LEVEL.value)
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+  params = parse_params()
+  logging.info('The actual training parameters are:\n%s', params.as_dict())
+  model_dir: str = os.path.join(
+      FLAGS.model_dir,
+      constants.TRIAL_PREFIX + hypertune_utils.get_trial_id_from_environment(),
+  )
+  logging.info('model_dir: %s', model_dir)
+
+  if 'train' in FLAGS.mode:
+    # Pure eval modes do not output yaml files. Otherwise continuous eval job
+    # may race against the train job for writing the same file.
+    train_utils.serialize_config(params, model_dir)
+
+  # Sets mixed_precision policy. Using 'mixed_float16' or 'mixed_bfloat16'
+  # can have significant impact on model speeds by utilizing float16 in case of
+  # GPUs, and bfloat16 in the case of TPUs. loss_scale takes effect only when
+  # dtype is float16
+  if params.runtime.mixed_precision_dtype:
+    performance.set_mixed_precision_policy(params.runtime.mixed_precision_dtype)
+  distribution_strategy = distribute_utils.get_distribution_strategy(
+      distribution_strategy=params.runtime.distribution_strategy,
+      all_reduce_alg=params.runtime.all_reduce_alg,
+      num_gpus=params.runtime.num_gpus,
+      tpu_address=params.runtime.tpu,
+  )
+
+  # Create task and run experiment.
+  with distribution_strategy.scope():
+    task = task_factory.get_task(params.task, logging_dir=model_dir)
+
+  train_lib.run_experiment(
+      distribution_strategy=distribution_strategy,
+      task=task,
+      mode=FLAGS.mode,
+      params=params,
+      model_dir=model_dir,
+  )
+
+  train_utils.save_gin_config(FLAGS.mode, model_dir)
+
+  eval_metric_name = constants.VIDEO_CLASSIFICATION_BEST_EVAL_METRIC
+
+  eval_filepath = os.path.join(
+      model_dir, constants.BEST_CKPT_DIRNAME, constants.BEST_CKPT_EVAL_FILENAME
+  )
+  logging.info('Load eval metrics from: %s.', eval_filepath)
+
+  with tf.io.gfile.GFile(eval_filepath, 'rb') as f:
+    eval_metric_results = json.load(f)
+  logging.info('eval metrics are: %s.', eval_metric_results)
+  if (
+      eval_metric_name in eval_metric_results
+      and constants.BEST_CKPT_STEP_NAME in eval_metric_results
+  ):
+    hp_metric = eval_metric_results[eval_metric_name]
+    hp_step = int(eval_metric_results[constants.BEST_CKPT_STEP_NAME])
+    hpt = hypertune.HyperTune()
+    hpt.report_hyperparameter_tuning_metric(
+        hyperparameter_metric_tag=constants.HP_METRIC_TAG,
+        metric_value=hp_metric,
+        global_step=hp_step,
+    )
+    logging.info(
+        'Send HP metric: %f and steps %d to hyperparameter tuning.',
+        hp_metric,
+        hp_step,
+    )
+  else:
+    logging.info(
+        'Either %s or %s is not included in the evaluation results: %s.',
+        eval_metric_name,
+        constants.BEST_CKPT_STEP_NAME,
+        eval_metric_results,
+    )
+
+
+if __name__ == '__main__':
+  tfm_flags.define_flags()
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/open_clip/dockerfile/serve.Dockerfile

@@ -0,0 +1,67 @@
+# Dockerfile for basic serving dockers for OpenCLIP.
+#
+# To build:
+# docker build -f model_oss/open_clip/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+# Switch to this base image for gpu serve.
+FROM pytorch/torchserve:0.7.1-gpu
+
+USER root
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="transformers_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Install libraries.
+RUN python3 -m pip install --upgrade pip
+RUN pip install torch==1.13.1
+RUN pip install open_clip_torch==2.20.0
+RUN pip install pillow==9.5.0
+RUN pip install google-cloud-storage==2.7.0
+
+# Copy model artifacts.
+COPY model_oss/open_clip/handler.py /home/model-server/handler.py
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server/
+
+# Create torchserve configuration file.
+RUN echo \
+    "default_response_timeout=1800\n" \
+    "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${infer_port}\n"  \
+    "management_address=http://0.0.0.0:${mng_port}" >> /home/model-server/config.properties
+
+# Expose ports.
+EXPOSE ${infer_port}
+EXPOSE ${mng_port}
+
+# Archive model artifacts and dependencies.
+# Do not set --model-file and --serialized-file because model and checkpoint will be dynamically loaded in handler.py.
+RUN torch-model-archiver \
+--model-name=${model_name} \
+--version=1.0 \
+--handler=/home/model-server/handler.py \
+--runtime=python3 \
+--export-path=/home/model-server/model-store \
+--archive-format=default \
+--force
+
+# Run Torchserve HTTP serve to respond to prediction requests.
+CMD ["torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${model_name}=${model_name}.mar", \
+     "--model-store", "/home/model-server/model-store"]

community-content/vertex_vision_model_garden/model_oss/open_clip/dockerfile/train.Dockerfile

@@ -0,0 +1,53 @@
+# Dockerfile for training dockers with OpenCLIP.
+#
+# To build:
+# docker build -f model_oss/open_clilp/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+RUN apt-get install -y --no-install-recommends apt-utils
+RUN apt-get install -y --no-install-recommends curl
+RUN apt-get install -y --no-install-recommends wget
+RUN apt-get install -y --no-install-recommends git
+RUN apt-get install -y --no-install-recommends jq
+RUN apt-get install -y --no-install-recommends gnupg
+RUN apt-get install -y --no-install-recommends build-essential
+
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Prepare artifacts.
+WORKDIR /workspace
+RUN git clone --branch main https://github.com/mlfoundations/open_clip.git
+WORKDIR ./open_clip
+RUN git reset --hard 67e5e5ec8741281eb9b30f640c26f91c666308b7
+
+# Install libraries.
+RUN pip install webdataset==0.2.5
+RUN pip install regex==2023.6.3
+RUN pip install ftfy==6.1.1
+RUN pip install pandas==2.0.3
+RUN pip install braceexpand==0.1.7
+RUN pip install huggingface_hub==0.16.4
+RUN pip install transformers==4.31.0
+RUN pip install timm==0.9.2
+RUN pip install fsspec==2023.6.0
+RUN pip install sentencepiece==0.1.99
+RUN pip install protobuf==3.20.3
+RUN pip install tensorboard==2.12.2
+
+# Switch work folder for training.
+WORKDIR ./src

community-content/vertex_vision_model_garden/model_oss/open_clip/handler.py

@@ -0,0 +1,142 @@
+"""Custom handler for OpenCLIP model."""
+
+# pylint:disable=g-importing-member
+import enum
+import logging
+import os
+from typing import Any, Dict, List
+
+import open_clip
+import torch
+from ts.torch_handler.base_handler import BaseHandler
+
+from google3.cloud.ml.applications.vision.model_garden.model_oss.util import constants
+from google3.cloud.ml.applications.vision.model_garden.model_oss.util import fileutils
+from google3.cloud.ml.applications.vision.model_garden.model_oss.util import image_format_converter
+
+
+@enum.unique
+class Precision(enum.Enum):
+  AMP = "amp"
+  AMP_BF16 = "amp_bf16"
+  AMP_BFLOAT16 = "amp_bfloat16"
+  BF16 = "bf16"
+  FP16 = "fp16"
+  PURE_BF16 = "pure_bf16"
+  PURE_FP16 = "pure_fp16"
+  FP32 = "fp32"
+
+
+# Supported checkpoint&model pairs:
+# https://github.com/mlfoundations/open_clip#pretrained-model-interface
+_DEFAULT_CHECKPOINT = "openai"
+_DEFAULT_MODEL = "RN50"
+_DEFAULT_PRECISION = Precision.AMP
+_ZERO_CLASSIFICATION = "zero-shot-image-classification"
+_FEATURE_EMBEDDING = "feature-embedding"
+_VALID_TASKS = frozenset([_ZERO_CLASSIFICATION, _FEATURE_EMBEDDING])
+
+_IMAGE_KEY = "image"
+_TEXT_KEY = "text"
+_IMAGE_FEATURES_KEY = "image_features"
+_TEXT_FEATURES_KEY = "text_features"
+
+
+class OpenclipHandler(BaseHandler):
+  """Custom handler for OpenCLIP."""
+
+  def initialize(self, context: Any):
+    """Custom initialize."""
+
+    properties = context.system_properties
+    self.map_location = (
+        "cuda"
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else "cpu"
+    )
+    self.device = torch.device(
+        self.map_location + ":" + str(properties.get("gpu_id"))
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else self.map_location
+    )
+    self.manifest = context.manifest
+
+    model_name = os.environ.get("MODEL", _DEFAULT_MODEL)
+    precision = os.environ.get("PRECISION", _DEFAULT_PRECISION)
+    checkpoint = os.environ.get("CHECKPOINT", _DEFAULT_CHECKPOINT)
+    self.task = os.environ.get("TASK", _FEATURE_EMBEDDING)
+    if self.task not in _VALID_TASKS:
+      raise ValueError(f"Invalid task: {self.task}.")
+    logging.info(
+        "Handler initializing task:%s, model:%s, precision:%s, checkpoint:%s",
+        self.task,
+        model_name,
+        precision,
+        checkpoint,
+    )
+
+    if checkpoint != _DEFAULT_CHECKPOINT:
+      local_fname = os.path.join(constants.LOCAL_MODEL_DIR, "model.pt")
+      fileutils.download_gcs_file_to_local(checkpoint, local_fname)
+      checkpoint = local_fname
+
+    self.model, _, self.preprocessor = open_clip.create_model_and_transforms(
+        model_name, pretrained=checkpoint, precision=precision
+    )
+    self.tokenizer = open_clip.get_tokenizer(model_name)
+
+    self.initialized = True
+
+  def preprocess(self, data: Any) -> List[Dict[str, Any]]:
+    """Preprocess input data."""
+    logging.info("preprocessing: %d instances received.", len(data))
+    processed_list = []
+    for item in data:
+      sample = {}
+      if _IMAGE_KEY in item:
+        sample[_IMAGE_KEY] = self.preprocessor(
+            image_format_converter.base64_to_image(item[_IMAGE_KEY])
+        ).unsqueeze(0)
+      if _TEXT_KEY in item:
+        sample[_TEXT_KEY] = self.tokenizer(item[_TEXT_KEY])
+      processed_list.append(sample)
+    return processed_list
+
+  def inference(
+      self, data: List[Dict[str, Any]], *args, **kwargs
+  ) -> List[Dict[str, Any]]:
+    feature_list = []
+    with torch.no_grad(), torch.cuda.amp.autocast():
+      for item in data:
+        sample = {}
+        if _IMAGE_KEY in item:
+          sample[_IMAGE_FEATURES_KEY] = self.model.encode_image(
+              item[_IMAGE_KEY]
+          )
+        if _TEXT_KEY in item:
+          sample[_TEXT_FEATURES_KEY] = self.model.encode_text(item[_TEXT_KEY])
+        feature_list.append(sample)
+    return feature_list
+
+  def postprocess(self, features: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Postprocess the image/text featreus for downstream task."""
+    preds = []
+    if self.task == _FEATURE_EMBEDDING:
+      for item in features:
+        preds.append({k: v.tolist() for k, v in item.items()})
+    elif self.task == _ZERO_CLASSIFICATION:
+      for item in features:
+        image_features = item.get(_IMAGE_FEATURES_KEY, None)
+        text_features = item.get(_TEXT_FEATURES_KEY, None)
+        if image_features is None or text_features is None:
+          raise ValueError(
+              "Missing input for {} task. {} received.".format(
+                  _ZERO_CLASSIFICATION, item.keys()
+              )
+          )
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+        preds.append(text_probs.tolist())
+
+    return preds

community-content/vertex_vision_model_garden/model_oss/peft/causal_language_modeling_lora.py

@@ -0,0 +1,142 @@
+"""Causal language modeling with LoRA models."""
+
+# pylint: disable=g-importing-member
+
+from datasets import load_dataset
+from peft import get_peft_model
+from peft import LoraConfig
+import torch
+from torch import nn
+import transformers
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import BitsAndBytesConfig
+from transformers import TrainingArguments
+from util import constants
+
+
+def finetune_causal_language_modeling(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    precision_mode: str = None,
+    lora_rank: int = 16,
+    lora_alpha: int = 32,
+    lora_dropout: float = 0.05,
+    warmup_steps: int = 10,
+    max_steps: int = 10,
+    learning_rate: float = 2e-4,
+    local_pretrained_model_id: str = None,
+) -> None:
+  """Finetunes causal language modelings."""
+  if precision_mode == constants.PRECISION_MODE_32:
+    model = AutoModelForCausalLM.from_pretrained(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        torch_dtype=torch.float32,
+        device_map="auto",
+    )
+  elif precision_mode == constants.PRECISION_MODE_16:
+    model = AutoModelForCausalLM.from_pretrained(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+  elif precision_mode == constants.PRECISION_MODE_8:
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=True, int8_threshold=0
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        quantization_config=quantization_config,
+    )
+  else:
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        quantization_config=quantization_config,
+    )
+
+  tokenizer = AutoTokenizer.from_pretrained(
+      local_pretrained_model_id
+      if local_pretrained_model_id
+      else pretrained_model_id
+  )
+  if "llama" in pretrained_model_id:
+    tokenizer.pad_token = "[PAD]"
+
+  for param in model.parameters():
+    # Freezes the model - train adapters later.
+    param.requires_grad = False
+    if param.ndim == 1:
+      # Casts the small parameters (e.g. layernorm) to fp32 for stability.
+      param.data = param.data.to(torch.float32)
+
+  # Reduces the number of stored activations.
+  model.gradient_checkpointing_enable()
+  model.enable_input_require_grads()
+
+  class CastOutputToFloat(nn.Sequential):
+
+    def forward(self, x):
+      return super().forward(x).to(torch.float32)
+
+  model.lm_head = CastOutputToFloat(model.lm_head)
+
+  config = LoraConfig(
+      r=lora_rank,
+      lora_alpha=lora_alpha,
+      target_modules=["q_proj", "v_proj"],
+      lora_dropout=lora_dropout,
+      bias="none",
+      task_type="CAUSAL_LM",
+  )
+
+  model = get_peft_model(model, config)
+  model.print_trainable_parameters()
+
+  data = load_dataset(dataset_name)
+  data = data.map(
+      lambda samples: tokenizer(samples["quote"]),
+      batched=True,
+  )
+
+  trainer = transformers.Trainer(
+      model=model,
+      train_dataset=data["train"],
+      args=TrainingArguments(
+          per_device_train_batch_size=4,
+          gradient_accumulation_steps=4,
+          warmup_steps=warmup_steps,
+          max_steps=max_steps,
+          learning_rate=learning_rate,
+          fp16=True,
+          logging_steps=1,
+          output_dir=output_dir,
+          ddp_find_unused_parameters=False,
+      ),
+      data_collator=transformers.DataCollatorForLanguageModeling(
+          tokenizer,
+          mlm=False,
+      ),
+  )
+  # Silence the warnings. Please re-enable for inference!
+  model.config.use_cache = False
+  trainer.train()
+
+  model.save_pretrained(output_dir)

community-content/vertex_vision_model_garden/model_oss/peft/config.properties

@@ -0,0 +1,21 @@
+number_of_netty_threads=32
+job_queue_size=1000
+model_store=/home/model-server/model-store
+workflow_store=/home/model-server/wf-store
+default_response_timeout=1800
+service_envelope=json
+inference_address=http://0.0.0.0:7080
+management_address=http://0.0.0.0:7081
+metrics_address=http://0.0.0.0:7082
+
+models={\
+  "peft_serving": {\
+    "1.0": {\
+        "defaultVersion": true,\
+        "marName": "peft_serving.mar",\
+        "minWorkers": 1,\
+        "maxWorkers": 1,\
+        "batchSize": 1\
+    }\
+  }\
+}

community-content/vertex_vision_model_garden/model_oss/peft/dockerfile/serve.Dockerfile

@@ -0,0 +1,107 @@
+# Dockerfile for PEFT Serving.
+#
+# To build:
+# docker build -f model_oss/peft/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM pytorch/torchserve:0.7.0-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="peft_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        git \
+        git-lfs
+RUN git lfs install
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install --upgrade torch==2.0.1
+RUN pip install torchvision==0.15.2
+RUN pip install tokenizers==0.13.3
+RUN pip install accelerate==0.21.0
+RUN pip install sentencepiece==0.1.99
+RUN pip install grpcio-status==1.33.2
+RUN pip install protobuf==3.19.6
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
+RUN pip install datasets==2.14.4
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.20
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
+RUN pip install scipy==1.10.1
+RUN pip install evaluate==0.4.0
+RUN pip install scikit-learn==1.2.2
+RUN pip install loralib==0.1.1
+RUN pip install bitsandbytes==0.39.0
+RUN pip install trl==0.4.4
+RUN pip install einops==0.6.1
+
+# Install diffusers from source.
+RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+WORKDIR /home/model-server
+
+# Install transformers from source.
+RUN git clone --depth 1 --branch v4.31.0 https://github.com/huggingface/transformers.git
+# The patch is used to change the transformers loading model behavior:
+# 1) For models on Huggingface hub: if the model has multiple shards, each shard
+#    will be downloaded separately and get deleted after loading to GPU.
+# 2) For models on local disk: if a model bin file is actually a text file
+#    recording a GCS path, the model file will be downloaded and get deleted
+#    after loading to GPU.
+COPY model_oss/peft/hf_transformers_lazy_download.patch /home/model-server/hf_transformers_lazy_download.patch
+WORKDIR transformers
+RUN git apply /home/model-server/hf_transformers_lazy_download.patch
+RUN pip install -e .
+WORKDIR /home/model-server
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY model_oss/peft/handler.py /home/model-server/handler.py
+COPY model_oss/peft/config.properties /home/model-server/config.properties
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server/
+
+# Expose ports.
+EXPOSE ${infer_port}
+EXPOSE ${mng_port}
+
+# Set environments.
+ENV TASK "causal-language-modeling-lora"
+ENV BASE_MODEL_ID "openlm-research/open_llama_7b"
+ENV PRECISION_LOADING_MODE "float16"
+ENV FINETUNED_LORA_MODEL_PATH ""
+
+
+# Archive model artifacts and dependencies.
+# Do not set --model-file and --serialized-file because model and checkpoint
+# will be dynamically loaded in handler.py.
+RUN torch-model-archiver \
+--model-name=${model_name} \
+--version=1.0 \
+--handler=/home/model-server/handler.py \
+--runtime=python3 \
+--export-path=/home/model-server/model-store \
+--archive-format=default \
+--force
+
+# Run Torchserve HTTP serve to respond to prediction requests.
+CMD ["torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${model_name}=${model_name}.mar", \
+     "--model-store", "/home/model-server/model-store"]

community-content/vertex_vision_model_garden/model_oss/peft/dockerfile/train.Dockerfile

@@ -0,0 +1,111 @@
+# Dockerfile for PEFT Training.
+#
+# To build:
+# docker build -f model_oss/peft/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+# Builds GPU docker image of PyTorch
+# Uses multi-staged approach to reduce size
+# Stage 1
+# Use base conda image to reduce time
+FROM continuumio/miniconda3:latest AS compile-image
+# Specify py version
+ENV PYTHON_VERSION=3.8
+# Install apt libs - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
+RUN apt-get update && \
+    apt-get install -y curl git wget software-properties-common git-lfs && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Install audio-related libraries
+RUN apt-get update && \
+    apt install -y ffmpeg
+
+RUN apt install -y libsndfile1-dev
+RUN git lfs install
+
+# Create our conda env - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
+RUN conda create --name peft python=${PYTHON_VERSION} ipython jupyter pip
+RUN python3 -m pip install --no-cache-dir --upgrade pip
+
+# Below is copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
+# We don't install pytorch here yet since CUDA isn't available
+# instead we use the direct torch wheel
+ENV PATH /opt/conda/envs/peft/bin:$PATH
+# Activate our bash shell
+RUN chsh -s /bin/bash
+SHELL ["/bin/bash", "-c"]
+# Activate the conda env and install transformers + accelerate from source
+RUN source activate peft
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/transformers
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/accelerate
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft#egg=peft[test]
+RUN python3 -m pip install --no-cache-dir bitsandbytes
+
+# Stage 2
+FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
+COPY --from=compile-image /opt/conda /opt/conda
+ENV PATH /opt/conda/bin:$PATH
+
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget vim && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN echo "source activate peft" >> ~/.profile
+
+# Install libraries.
+RUN pip install --upgrade torch==2.0.1
+RUN pip install torchvision==0.15.2
+RUN pip install git+https://github.com/huggingface/transformers@de9255de27abfcae4a1f816b904915f0b1e23cd9
+RUN pip install transformers -U
+RUN pip install accelerate==0.21.0
+RUN pip install sentencepiece==0.1.99
+RUN pip install grpcio-status==1.33.2
+RUN pip install protobuf==3.19.6
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
+RUN pip install datasets==2.9.0
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.20
+RUN pip install Jinja2==3.1.2
+RUN pip install ftfy==6.1.1
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install tensorboard==2.12.0
+RUN pip install scipy==1.10.1
+RUN pip install evaluate==0.4.0
+RUN pip install scikit-learn==1.2.2
+RUN pip install loralib==0.1.1
+RUN pip install bitsandbytes==0.39.0
+RUN pip install trl==0.4.4
+RUN pip install einops==0.6.1
+RUN pip install google-cloud-storage==2.7.0
+
+RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+
+# Switch to diffusers examples folder.
+WORKDIR examples
+
+# NOTE: use 'sed' to modify train_text_to_image_lora.py to
+#       fix the bug for accelerator.
+RUN sed -i \
+"s#logging_dir=logging_dir#project_dir=logging_dir#g" \
+text_to_image/train_text_to_image_lora.py
+
+# Config accelerate.
+RUN mkdir -p ./vertex_vision_model_garden_peft/
+COPY model_oss/peft/train.sh ./vertex_vision_model_garden_peft/train.sh
+COPY model_oss/peft/*.py     ./vertex_vision_model_garden_peft/
+COPY model_oss/util /diffusers/examples/util
+ENV PYTHONPATH /diffusers/examples/
+
+# Generate accelerate config at the beginning of docker run.
+ENTRYPOINT ["python3", "vertex_vision_model_garden_peft/main.py"]

community-content/vertex_vision_model_garden/model_oss/peft/handler.py

@@ -0,0 +1,250 @@
+"""Custom handler for huggingface/peft models."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=logging-fstring-interpolation
+
+import logging
+import os
+from typing import Any, List
+
+from absl import logging
+from diffusers import DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline
+from peft import PeftModel
+from PIL import Image
+import torch
+import transformers
+from transformers import AutoModelForCausalLM
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+from transformers import BitsAndBytesConfig
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import constants
+from util import fileutils
+from util import image_format_converter
+
+# Tasks
+TEXT_TO_IMAGE_LORA = "text-to-image-lora"
+SEQUENCE_CLASSIFICATION_LORA = "sequence-classification-lora"
+CAUSAL_LANGUAGE_MODELING_LORA = "causal-language-modeling-lora"
+INSTRUCT_LORA = "instruct-lora"
+
+# Inference parameters.
+_NUM_INFERENCE_STEPS = 25
+_MAX_LENGTH_DEFAULT = 200
+_TOP_K_DEFAULT = 10
+
+
+class PeftHandler(BaseHandler):
+  """Custom handler for Peft models."""
+
+  def initialize(self, context: Any):
+    """Initializes the handler."""
+    logging.info("Start to initialize the PEFT handler.")
+    properties = context.system_properties
+    self.map_location = (
+        "cuda"
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else "cpu"
+    )
+
+    self.device = torch.device(
+        self.map_location + ":" + str(properties.get("gpu_id"))
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else self.map_location
+    )
+    self.manifest = context.manifest
+    self.precision_mode = os.environ.get(
+        "PRECISION_LOADING_MODE", constants.PRECISION_MODE_16
+    )
+    self.task = os.environ.get("TASK", CAUSAL_LANGUAGE_MODELING_LORA)
+    self.base_model_id = os.environ.get(
+        "BASE_MODEL_ID", "openlm-research/open_llama_7b"
+    )
+    if fileutils.is_gcs_path(self.base_model_id):
+      fileutils.download_gcs_dir_to_local(
+          self.base_model_id,
+          constants.LOCAL_BASE_MODEL_DIR,
+          skip_hf_model_bin=True,
+      )
+      self.base_model_id = constants.LOCAL_BASE_MODEL_DIR
+    self.finetuned_lora_model_path = os.environ.get(
+        "FINETUNED_LORA_MODEL_PATH", ""
+    )
+    if fileutils.is_gcs_path(self.finetuned_lora_model_path):
+      fileutils.download_gcs_dir_to_local(
+          self.finetuned_lora_model_path, constants.LOCAL_MODEL_DIR
+      )
+      self.finetuned_lora_model_path = constants.LOCAL_MODEL_DIR
+
+    logging.info(
+        f"Using task:{self.task}, base model:{self.base_model_id}, lora model:"
+        f" {self.finetuned_lora_model_path}, and precision"
+        f" {self.precision_mode}."
+    )
+
+    self.pipeline = None
+    self.model = None
+    self.tokenizer = None
+    if self.task == TEXT_TO_IMAGE_LORA:
+      pipeline = StableDiffusionPipeline.from_pretrained(
+          self.base_model_id, torch_dtype=torch.float16
+      )
+      logging.debug("Initialized the base model for text to image.")
+      pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      logging.debug("Initialized the scheduler for text to image.")
+      if self.finetuned_lora_model_path:
+        pipeline.unet.load_attn_procs(self.finetuned_lora_model_path)
+      logging.debug("Initialized the LoRA model for text to image.")
+      # This is to reduce GPU memory requirements.
+      pipeline.enable_xformers_memory_efficient_attention()
+      pipeline = pipeline.to(self.map_location)
+      # Reduces memory footprint.
+      pipeline.enable_attention_slicing()
+      self.pipeline = pipeline
+      logging.info("Initialized the text to image pipelines.")
+    elif self.task == SEQUENCE_CLASSIFICATION_LORA:
+      tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
+      logging.debug("Initialized the tokenizer for sequence classification.")
+      model = AutoModelForSequenceClassification.from_pretrained(
+          self.base_model_id, torch_dtype=torch.float16
+      )
+      logging.debug("Initialized the base model for sequence classification.")
+      if self.finetuned_lora_model_path:
+        model = PeftModel.from_pretrained(model, self.finetuned_lora_model_path)
+        logging.debug("Initialized the LoRA model for sequence classification.")
+      model.to(self.map_location)
+      self.model = model
+      self.tokenizer = tokenizer
+    elif (
+        self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
+    ):
+      tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
+      logging.debug("Initialized the tokenizer.")
+      if self.task == CAUSAL_LANGUAGE_MODELING_LORA:
+        if self.precision_mode == constants.PRECISION_MODE_32:
+          model = AutoModelForCausalLM.from_pretrained(
+              self.base_model_id,
+              return_dict=True,
+              torch_dtype=torch.float32,
+              device_map="auto",
+          )
+        elif self.precision_mode == constants.PRECISION_MODE_16:
+          model = AutoModelForCausalLM.from_pretrained(
+              self.base_model_id,
+              return_dict=True,
+              torch_dtype=torch.bfloat16,
+              device_map="auto",
+          )
+        elif self.precision_mode == constants.PRECISION_MODE_8:
+          quantization_config = BitsAndBytesConfig(
+              load_in_8bit=True, int8_threshold=0
+          )
+          model = AutoModelForCausalLM.from_pretrained(
+              self.base_model_id,
+              return_dict=True,
+              torch_dtype=torch.float16,
+              device_map="auto",
+              quantization_config=quantization_config,
+          )
+        else:
+          quantization_config = BitsAndBytesConfig(
+              load_in_4bit=True,
+              bnb_4bit_quant_type="nf4",
+              bnb_4bit_compute_dtype=torch.bfloat16,
+          )
+          model = AutoModelForCausalLM.from_pretrained(
+              self.base_model_id,
+              return_dict=True,
+              device_map="auto",
+              torch_dtype=torch.bfloat16,
+              quantization_config=quantization_config,
+          )
+      else:
+        model = AutoModelForCausalLM.from_pretrained(
+            self.base_model_id,
+            return_dict=True,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map="auto",
+        )
+      logging.debug("Initialized the base model.")
+      if self.finetuned_lora_model_path:
+        model = PeftModel.from_pretrained(model, self.finetuned_lora_model_path)
+        logging.debug("Initialized the LoRA model.")
+      pipeline = transformers.pipeline(
+          "text-generation",
+          model=model,
+          tokenizer=tokenizer,
+      )
+      self.tokenizer = tokenizer
+      self.pipeline = pipeline
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+
+    self.initialized = True
+    logging.info("The PEFT handler was initialized.")
+
+  def preprocess(self, data: Any) -> Any:
+    """Preprocesses input data."""
+    # Assumes that the parameters are same in one request. We parse the
+    # parameters from the first instance for all instances in one request.
+    max_length = _MAX_LENGTH_DEFAULT
+    top_k = _TOP_K_DEFAULT
+
+    prompts = [item["prompt"] for item in data]
+    if "max_length" in data[0]:
+      max_length = data[0]["max_length"]
+    if "top_k" in data[0]:
+      top_k = data[0]["top_k"]
+
+    return prompts, max_length, top_k
+
+  def inference(self, data: Any, *args, **kwargs) -> List[Image.Image]:
+    """Runs the inference."""
+    prompts, max_length, top_k = data
+    logging.debug(
+        f"Inference prompts={prompts}, max_length={max_length}, top_k={top_k}."
+    )
+    if self.task == TEXT_TO_IMAGE_LORA:
+      predicted_results = self.pipeline(
+          prompt=prompts, num_inference_steps=_NUM_INFERENCE_STEPS
+      ).images
+    elif self.task == SEQUENCE_CLASSIFICATION_LORA:
+      encoded_input = self.tokenizer(prompts, return_tensors="pt")
+      encoded_input.to(self.map_location)
+      with torch.no_grad():
+        outputs = self.model(**encoded_input)
+      predictions = outputs.logits.argmax(dim=-1)
+      predicted_results = predictions.tolist()
+    elif (
+        self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
+    ):
+      predicted_results = self.pipeline(
+          prompts,
+          max_length=max_length,
+          do_sample=True,
+          top_k=top_k,
+          num_return_sequences=1,
+          eos_token_id=self.tokenizer.eos_token_id,
+      )
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+    return predicted_results
+
+  def postprocess(self, data: Any) -> List[str]:
+    """Postprocesses output data."""
+    if self.task == TEXT_TO_IMAGE_LORA:
+      # Converts the images to base64 string.
+      outputs = [
+          image_format_converter.image_to_base64(image) for image in data
+      ]
+    else:
+      outputs = data
+    return outputs
+
+
+# pylint: enable=logging-fstring-interpolation

community-content/vertex_vision_model_garden/model_oss/peft/hf_transformers_lazy_download.patch

@@ -0,0 +1,131 @@
+diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
+index 45459ed..32527f4 100644
+--- a/src/transformers/modeling_utils.py
++++ b/src/transformers/modeling_utils.py
+@@ -32,6 +32,8 @@ import torch
+ from packaging import version
+ from torch import Tensor, nn
+ from torch.nn import CrossEntropyLoss
++from huggingface_hub import hf_hub_download
++from google.cloud import storage
+
+ from .activations import get_activation
+ from .configuration_utils import PretrainedConfig
+@@ -442,6 +444,29 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
+     """
+     Reads a PyTorch checkpoint file, returning properly formatted errors if they arise.
+     """
++    delete_download = False
++    tmp_dir = "/tmp/model"
++    os.makedirs(tmp_dir, exist_ok=True)
++    if isinstance(checkpoint_file, dict):
++        # Download model file from huggingface
++        print(f"==> Download model from HF: {checkpoint_file}")
++        checkpoint_file = hf_hub_download(
++            local_dir=tmp_dir, local_dir_use_symlinks=False, force_download=True, resume_download=True, **checkpoint_file)
++        delete_download = True
++    else:
++        with open(checkpoint_file, "rb") as f:
++            is_gcs_file = (f.read(2) == b"gs")
++        if is_gcs_file:
++            # Download model file from GCS
++            with open(checkpoint_file, "r") as f:
++                gcs_file = f.read()
++            checkpoint_file = os.path.join(tmp_dir, gcs_file.split("/")[-1])
++            print(f"==> Download model from GCS: {gcs_file} to: {checkpoint_file}")
++            client = storage.Client()
++            with open(checkpoint_file, 'wb') as f:
++                client.download_blob_to_file(gcs_file, f)
++            delete_download = True
++
+     if checkpoint_file.endswith(".safetensors") and is_safetensors_available():
+         # Check format of the archive
+         with safe_open(checkpoint_file, framework="pt") as f:
+@@ -455,9 +480,9 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
+             raise NotImplementedError(
+                 f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet."
+             )
+-        return safe_load_file(checkpoint_file)
++        state_dict = safe_load_file(checkpoint_file)
+     try:
+-        return torch.load(checkpoint_file, map_location="cpu")
++        state_dict = torch.load(checkpoint_file, map_location="cpu")
+     except Exception as e:
+         try:
+             with open(checkpoint_file) as f:
+@@ -478,6 +503,10 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
+                 f"at '{checkpoint_file}'. "
+                 "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
+             )
++    if delete_download:
++        print(f"==> Delete downloaded model: {checkpoint_file}")
++        os.remove(checkpoint_file)
++    return state_dict
+
+
+ def set_initialized_submodules(model, state_dict_keys):
+@@ -3179,7 +3208,10 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
+             return mismatched_keys
+
+         if resolved_archive_file is not None:
+-            folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
++            if isinstance(resolved_archive_file, str):
++                folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
++            else:
++                folder = None
+         else:
+             folder = None
+         if device_map is not None and is_safetensors:
+diff --git a/src/transformers/utils/hub.py b/src/transformers/utils/hub.py
+index ffed743..4b15770 100644
+--- a/src/transformers/utils/hub.py
++++ b/src/transformers/utils/hub.py
+@@ -414,20 +414,34 @@ def cached_file(
+     user_agent = http_user_agent(user_agent)
+     try:
+         # Load from URL or cache if already cached
+-        resolved_file = hf_hub_download(
+-            path_or_repo_id,
+-            filename,
+-            subfolder=None if len(subfolder) == 0 else subfolder,
+-            repo_type=repo_type,
+-            revision=revision,
+-            cache_dir=cache_dir,
+-            user_agent=user_agent,
+-            force_download=force_download,
+-            proxies=proxies,
+-            resume_download=resume_download,
+-            use_auth_token=use_auth_token,
+-            local_files_only=local_files_only,
+-        )
++        if filename.endswith(".bin"):
++            # NOTE: To save disk we do not download bin file eagerly. Do not support safetensors.
++            resolved_file = dict(
++                repo_id=path_or_repo_id,
++                filename=filename,
++                subfolder=None if len(subfolder) == 0 else subfolder,
++                repo_type=repo_type,
++                revision=revision,
++                user_agent=user_agent,
++                proxies=proxies,
++                use_auth_token=use_auth_token,
++            )
++            print(f"--> Apply lazy download to bin file: {resolved_file}")
++        else:
++            resolved_file = hf_hub_download(
++                path_or_repo_id,
++                filename,
++                subfolder=None if len(subfolder) == 0 else subfolder,
++                repo_type=repo_type,
++                revision=revision,
++                cache_dir=cache_dir,
++                user_agent=user_agent,
++                force_download=force_download,
++                proxies=proxies,
++                resume_download=resume_download,
++                use_auth_token=use_auth_token,
++                local_files_only=local_files_only,
++            )
+
+     except RepositoryNotFoundError:
+         raise EnvironmentError(

community-content/vertex_vision_model_garden/model_oss/peft/instruct_lora.py

@@ -0,0 +1,97 @@
+"""Instruct/Chat with LoRA models."""
+
+# pylint: disable=g-importing-member
+from datasets import load_dataset
+from peft import LoraConfig
+import torch
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import BitsAndBytesConfig
+from transformers import TrainingArguments
+from trl import SFTTrainer
+
+
+def finetune_instruct(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    lora_rank: int = 64,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.1,
+    warmup_ratio: int = 0.03,
+    max_steps: int = 10,
+    max_seq_length: int = 512,
+    learning_rate: float = 2e-4,
+) -> None:
+  """Finetunes instruct."""
+  dataset = load_dataset(dataset_name, split="train")
+
+  bnb_config = BitsAndBytesConfig(
+      load_in_4bit=True,
+      bnb_4bit_quant_type="nf4",
+      bnb_4bit_compute_dtype=torch.float16,
+  )
+
+  model = AutoModelForCausalLM.from_pretrained(
+      pretrained_model_id,
+      quantization_config=bnb_config,
+      trust_remote_code=True,
+  )
+  model.config.use_cache = False
+
+  tokenizer = AutoTokenizer.from_pretrained(
+      pretrained_model_id, trust_remote_code=True
+  )
+  tokenizer.pad_token = tokenizer.eos_token
+
+  peft_config = LoraConfig(
+      lora_alpha=lora_alpha,
+      lora_dropout=lora_dropout,
+      r=lora_rank,
+      bias="none",
+      task_type="CAUSAL_LM",
+      target_modules=[
+          "query_key_value",
+          "dense",
+          "dense_h_to_4h",
+          "dense_4h_to_h",
+      ],
+  )
+
+  per_device_train_batch_size = 4
+  gradient_accumulation_steps = 4
+  optim = "paged_adamw_32bit"
+  save_steps = 10
+  logging_steps = 10
+  max_grad_norm = 0.3
+  lr_scheduler_type = "constant"
+
+  training_arguments = TrainingArguments(
+      output_dir=output_dir,
+      per_device_train_batch_size=per_device_train_batch_size,
+      gradient_accumulation_steps=gradient_accumulation_steps,
+      optim=optim,
+      save_steps=save_steps,
+      logging_steps=logging_steps,
+      learning_rate=learning_rate,
+      fp16=True,
+      max_grad_norm=max_grad_norm,
+      max_steps=max_steps,
+      warmup_ratio=warmup_ratio,
+      group_by_length=True,
+      lr_scheduler_type=lr_scheduler_type,
+  )
+
+  trainer = SFTTrainer(
+      model=model,
+      train_dataset=dataset,
+      peft_config=peft_config,
+      dataset_text_field="text",
+      max_seq_length=max_seq_length,
+      tokenizer=tokenizer,
+      args=training_arguments,
+  )
+  for name, module in trainer.model.named_modules():
+    if "norm" in name:
+      module = module.to(torch.float32)
+  trainer.train()

community-content/vertex_vision_model_garden/model_oss/peft/main.py

@@ -0,0 +1,177 @@
+"""Main function to start PEFT finetuning."""
+import subprocess
+
+from absl import app
+from absl import flags
+from absl import logging
+
+from peft import causal_language_modeling_lora
+from peft import instruct_lora
+from peft import sequence_classification_lora
+from util import constants
+from util import fileutils
+
+_TASK = flags.DEFINE_string(
+    'task',
+    constants.CAUSAL_LANGUAGE_MODELING_LORA,
+    'The supported PEFT tasks.',
+)
+
+_PRETRAINED_MODEL_ID = flags.DEFINE_string(
+    'pretrained_model_id',
+    None,
+    'The pretrained model id. Supported models can be causal language modeling'
+    ' models from https://github.com/huggingface/peft/tree/main.',
+    required=True,
+)
+
+_DATASET_NAME = flags.DEFINE_string(
+    'dataset_name',
+    None,
+    'The dataset name in huggingface.',
+    required=True,
+)
+
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'The output directory.',
+    required=True,
+)
+
+_PRECISION_MODE = flags.DEFINE_string(
+    'precision_mode',
+    constants.PRECISION_MODE_16,
+    'Supported finetuning precision_modes are `{}` and `{}`.'.format(
+        constants.PRECISION_MODE_8, constants.PRECISION_MODE_16
+    ),
+)
+
+_LORA_RANK = flags.DEFINE_integer(
+    'lora_rank',
+    16,
+    'The rank of the update matrices, expressed in int. Lower rank results in'
+    ' smaller update matrices with fewer trainable parameters, referring to'
+    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
+)
+
+_LORA_ALPHA = flags.DEFINE_integer(
+    'lora_alpha',
+    32,
+    'LoRA scaling factor, referring to'
+    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
+)
+
+_LORA_DROPOUT = flags.DEFINE_float(
+    'lora_dropout',
+    0.05,
+    'dropout probability of the LoRA layers, referring to'
+    ' https://huggingface.co/docs/peft/task_guides/token-classification-lora.',
+)
+
+_WARMUP_STEPS = flags.DEFINE_integer(
+    'warmup_steps',
+    10,
+    'Number of steps for the warmup in the learning rate scheduler.',
+)
+
+_WARMUP_RATIO = flags.DEFINE_float(
+    'warmup_ratio',
+    0.03,
+    'The warmup ratio in the learning rate scheduler.',
+)
+
+_MAX_STEPS = flags.DEFINE_integer(
+    'max_steps',
+    10,
+    'Total number of training steps.',
+)
+
+_MAX_SEQ_LENGTH = flags.DEFINE_integer(
+    'max_seq_length',
+    512,
+    'The maximum sequence length.',
+)
+
+_NUM_EPOCHS = flags.DEFINE_integer(
+    'num_epochs',
+    20,
+    'The number of training epochs.',
+)
+
+_BATCH_SIZE = flags.DEFINE_integer(
+    'batch_size',
+    32,
+    'The batch size.',
+)
+
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate',
+    2e-4,
+    'The learning rate after the potential warmup period.',
+)
+
+
+def main(_) -> None:
+  task = _TASK.value
+  pretrained_model_id = _PRETRAINED_MODEL_ID.value
+  local_pretrained_model_id = None
+  if pretrained_model_id.startswith(constants.GCS_URI_PREFIX):
+    logging.info(
+        'Start to copy pretrained models locally: %s.', pretrained_model_id
+    )
+    fileutils.download_gcs_dir_to_local(
+        pretrained_model_id, constants.LOCAL_BASE_MODEL_DIR
+    )
+    local_pretrained_model_id = constants.LOCAL_BASE_MODEL_DIR
+    logging.info(
+        'Finished copying pretrained models locally to: %s.',
+        local_pretrained_model_id,
+    )
+  if task == constants.TEXT_TO_IMAGE_LORA:
+    subprocess.run(['/bin/bash', 'train.sh'], check=True)
+  elif task == constants.SEQUENCE_CLASSIFICATION_LORA:
+    sequence_classification_lora.finetune_sequence_classification(
+        pretrained_model_id=pretrained_model_id,
+        dataset_name=_DATASET_NAME.value,
+        output_dir=_OUTPUT_DIR.value,
+        lora_rank=_LORA_RANK.value,
+        lora_alpha=_LORA_ALPHA.value,
+        lora_dropout=_LORA_DROPOUT.value,
+        num_epochs=_NUM_EPOCHS.value,
+        batch_size=_BATCH_SIZE.value,
+        learning_rate=_LEARNING_RATE.value,
+    )
+  elif task == constants.CAUSAL_LANGUAGE_MODELING_LORA:
+    causal_language_modeling_lora.finetune_causal_language_modeling(
+        pretrained_model_id=pretrained_model_id,
+        dataset_name=_DATASET_NAME.value,
+        output_dir=_OUTPUT_DIR.value,
+        precision_mode=_PRECISION_MODE.value,
+        lora_rank=_LORA_RANK.value,
+        lora_alpha=_LORA_ALPHA.value,
+        lora_dropout=_LORA_DROPOUT.value,
+        warmup_steps=_WARMUP_STEPS.value,
+        max_steps=_MAX_STEPS.value,
+        learning_rate=_LEARNING_RATE.value,
+        local_pretrained_model_id=local_pretrained_model_id,
+    )
+  elif task == constants.INSTRUCT_LORA:
+    instruct_lora.finetune_instruct(
+        pretrained_model_id=pretrained_model_id,
+        dataset_name=_DATASET_NAME.value,
+        output_dir=_OUTPUT_DIR.value,
+        lora_rank=_LORA_RANK.value,
+        lora_alpha=_LORA_ALPHA.value,
+        lora_dropout=_LORA_DROPOUT.value,
+        warmup_ratio=_WARMUP_RATIO.value,
+        max_steps=_MAX_STEPS.value,
+        max_seq_length=_MAX_SEQ_LENGTH.value,
+        learning_rate=_LEARNING_RATE.value,
+    )
+  else:
+    raise ValueError('The task {} is not supported.'.format(task))
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/peft/sequence_classification_lora.py

@@ -0,0 +1,133 @@
+"""Sequence classification with LoRA models."""
+
+# pylint: disable=g-importing-member
+
+from datasets import load_dataset
+import evaluate
+from peft import get_peft_model
+from peft import LoraConfig
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+from transformers import get_linear_schedule_with_warmup
+
+
+def finetune_sequence_classification(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    lora_rank: int = 8,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.1,
+    num_epochs: int = 20,
+    batch_size: int = 32,
+    learning_rate: float = 3e-4,
+) -> None:
+  """Finetunes sequence classification."""
+  task = "mrpc"
+  device = "cuda"
+
+  peft_config = LoraConfig(
+      task_type="SEQ_CLS",
+      inference_mode=False,
+      r=lora_rank,
+      lora_alpha=lora_alpha,
+      lora_dropout=lora_dropout,
+  )
+  if any(k in pretrained_model_id for k in ("gpt", "opt", "bloom")):
+    padding_side = "left"
+  else:
+    padding_side = "right"
+
+  tokenizer = AutoTokenizer.from_pretrained(
+      pretrained_model_id, padding_side=padding_side
+  )
+  if getattr(tokenizer, "pad_token_id") is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
+  datasets = load_dataset(dataset_name, task)
+  metric = evaluate.load(dataset_name, task)
+
+  def tokenize_function(examples):
+    # max_length=None => use the model max length (it's actually the default)
+    outputs = tokenizer(
+        examples["sentence1"],
+        examples["sentence2"],
+        truncation=True,
+        max_length=None,
+    )
+    return outputs
+
+  tokenized_datasets = datasets.map(
+      tokenize_function,
+      batched=True,
+      remove_columns=["idx", "sentence1", "sentence2"],
+  )
+
+  # We also rename the 'label' column to 'labels' which is the expected name for
+  # labels by the models of the transformers library.
+  tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+  def collate_fn(examples):
+    return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+  # Instantiate dataloaders.
+  train_dataloader = DataLoader(
+      tokenized_datasets["train"],
+      shuffle=True,
+      collate_fn=collate_fn,
+      batch_size=batch_size,
+  )
+  eval_dataloader = DataLoader(
+      tokenized_datasets["validation"],
+      shuffle=False,
+      collate_fn=collate_fn,
+      batch_size=batch_size,
+  )
+
+  model = AutoModelForSequenceClassification.from_pretrained(
+      pretrained_model_id, return_dict=True
+  )
+  model = get_peft_model(model, peft_config)
+  model.print_trainable_parameters()
+
+  optimizer = AdamW(params=model.parameters(), lr=learning_rate)
+
+  # Instantiate scheduler
+  lr_scheduler = get_linear_schedule_with_warmup(
+      optimizer=optimizer,
+      num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
+      num_training_steps=(len(train_dataloader) * num_epochs),
+  )
+
+  model.to(device)
+  for epoch in range(num_epochs):
+    model.train()
+    for _, batch in enumerate(tqdm(train_dataloader)):
+      batch.to(device)
+      outputs = model(**batch)
+      loss = outputs.loss
+      loss.backward()
+      optimizer.step()
+      lr_scheduler.step()
+      optimizer.zero_grad()
+
+    model.eval()
+    for _, batch in enumerate(tqdm(eval_dataloader)):
+      batch.to(device)
+      with torch.no_grad():
+        outputs = model(**batch)
+      predictions = outputs.logits.argmax(dim=-1)
+      references = batch["labels"]
+      metric.add_batch(
+          predictions=predictions,
+          references=references,
+      )
+
+    eval_metric = metric.compute()
+    print(f"epoch {epoch}:", eval_metric)
+
+  model.save_pretrained(output_dir)

community-content/vertex_vision_model_garden/model_oss/peft/train.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+# Setup accelerate config before running trainer.
+python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+
+accelerate launch "$@"

community-content/vertex_vision_model_garden/model_oss/pic2word/dockerfile/serve.Dockerfile

@@ -0,0 +1,115 @@
+FROM pytorch/torchserve:0.7.1-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="pic2word"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Copy license.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        wget
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install dependencies.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install open_clip_torch==2.20.0
+RUN pip install numpy==1.22.0
+RUN pip install scikit-image==0.21.0
+RUN pip install scikit-learn==1.0.2
+RUN pip install torch==2.0.0
+RUN pip install torchvision==0.15.2
+RUN pip install tensorboard==2.13.0
+RUN pip install ase==3.21.1
+RUN pip install braceexpand==0.1.7
+RUN pip install cached-property==1.5.2
+RUN pip install configparser==5.0.2
+RUN pip install cycler==0.10.0
+RUN pip install decorator==4.4.2
+RUN pip install docker-pycreds==0.4.0
+RUN pip install gitdb==4.0.7
+RUN pip install gitpython==3.1.30
+RUN pip install googledrivedownloader==0.4
+RUN pip install h5py==3.1.0
+RUN pip install isodate==0.6.0
+RUN pip install jinja2==3.0.1
+RUN pip install kiwisolver==1.3.1
+RUN pip install littleutils==0.2.2
+RUN pip install llvmlite==0.36.0
+RUN pip install markupsafe==2.0.1
+RUN pip install matplotlib==3.3.4
+RUN pip install networkx==2.5.1
+RUN pip install numba==0.53.1
+RUN pip install ogb==1.3.1
+RUN pip install outdated==0.2.1
+RUN pip install pathtools==0.1.2
+RUN pip install promise==2.3
+RUN pip install psutil==5.8.0
+RUN pip install pyarrow==4.0.0
+RUN pip install pyparsing==2.4.7
+RUN pip install python-louvain==0.15
+RUN pip install pyyaml==5.4.1
+RUN pip install rdflib==5.0.0
+RUN pip install sentry-sdk==1.14.0
+RUN pip install shortuuid==1.0.1
+RUN pip install sklearn==0.0
+RUN pip install smmap==4.0.0
+RUN pip install subprocess32==3.5.4
+RUN pip install torch-geometric==1.7.0
+RUN pip install wandb==0.10.30
+RUN pip install wilds==1.1.0
+RUN pip install ftfy==6.1.1
+RUN pip install regex==2023.6.3
+RUN pip install webdataset==0.2.48
+RUN pip install requests==2.31.0
+RUN pip install hydra-core==1.3.2
+RUN pip install omegaconf==2.3.0
+RUN pip install fairseq==0.10.0
+RUN pip install bitarray==2.7.6
+
+# Get 'composed_image_retrieval' repository from github.
+RUN git clone https://github.com/google-research/composed_image_retrieval
+# Set workdir to composed_image_retrieval.
+WORKDIR ./composed_image_retrieval
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard 8c053297c2fae9cd17ddcded48445a4f47208dbd
+
+# Fix issue introduced by installing composed_image_retrieval
+# https://github.com/huggingface/transformers/issues/8638#issuecomment-790772391
+RUN pip uninstall dataclasses -y
+
+# Copy model artifacts.
+COPY model_oss/pic2word/handler.py /home/model-server/handler.py
+
+# Create torchserve configuration file.
+RUN echo \
+    "default_response_timeout=1800\n" \
+    "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${infer_port}\n"  \
+    "management_address=http://0.0.0.0:${mng_port}" >> /home/model-server/config.properties
+
+# Expose ports.
+EXPOSE ${infer_port}
+EXPOSE ${mng_port}
+
+# Archive model artifacts and dependencies.
+# Do not set --model-file and --serialized-file because model and checkpoint
+# will be dynamically loaded in handler.py.
+RUN torch-model-archiver \
+--model-name=${model_name} \
+--version=1.0 \
+--handler=/home/model-server/handler.py \
+--runtime=python3 \
+--export-path=/home/model-server/model-store \
+--archive-format=default \
+--force
+
+# Run Torchserve HTTP serve to respond to prediction requests.
+CMD ["torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${model_name}=${model_name}.mar", \
+     "--model-store", "/home/model-server/model-store"]

community-content/vertex_vision_model_garden/model_oss/pic2word/handler.py

@@ -0,0 +1,167 @@
+"""Custom handler for Pic2Word."""
+
+from argparse import Namespace  # pylint: disable=g-importing-member
+import os
+from typing import Any
+
+from absl import logging
+from data import CustomFolder
+from eval_utils import visualize_results
+from model.clip import load
+from model.model import convert_weights
+from model.model import IM2TEXT
+from params import get_project_root
+import torch
+from torch.utils.data import DataLoader
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import fileutils
+
+# The COCO dataset is stored in a publicly accessible bucket.
+_COCO_STORAGE_DIR = "gs://pic2word-bucket/data/coco/"
+_COCO_LOCAL_DIR = "/home/model-server/composed_image_retrieval/data/coco/"
+_COCO_VAL2017_PATH = "coco/val2017"
+_COCO_DATASET_NAME = "coco"
+_MODEL_NAME = "ViT-L/14"
+_LOCAL_QUERY_PATH = "./query/"
+_IMAGE_OUTPUT_LOCAL_DIR = "demo_out/images"
+_OUTPUT_LOCAL_DIR = "/demo_out/"
+_DATA_DIR = "data"
+_CHECKPOINT_DIR = "checkpoint/pic2word_model.pt"
+_REQUEST_PROMPTS = "prompts"
+_REQUEST_OUTPUT_STORAGE_DIR = "output_storage_dir"
+_REQUEST_IMAGE_PATH = "image_path"
+_REQUEST_IMAGE_FILE_NAME = "image_file_name"
+_RESPONSE_MSG = "Successfully retrieved images."
+
+
+class ModelHandler(BaseHandler):
+  """A custom model handler implementation."""
+
+  def __init__(self):
+    self.initialized = False
+    self.gpu = 0
+    self.model = None
+    self.dataloader = None
+    self.prompt = None
+    self.output_storage_dir = None
+
+  def initialize(self, context: Any):
+    """Initialize."""
+    logging.info("Initializing pic2word.")
+
+    # Download COCO dataset. The model looks for this folder specifically
+    # during image retrieval to generate a response for each request.
+    # This is a publicly accessible bucket.
+    fileutils.download_gcs_dir_to_local(
+        _COCO_STORAGE_DIR,
+        _COCO_LOCAL_DIR,
+    )
+
+    # Load the model.
+
+    self.initialized = True
+
+    torch.cuda.set_device(self.gpu)
+    model, _, preprocess_val = load(_MODEL_NAME, jit=False)
+
+    img2text = IM2TEXT(
+        embed_dim=model.embed_dim,
+        output_dim=model.token_embedding.weight.shape[1],
+    )
+
+    model.cuda(self.gpu)
+    img2text.cuda(self.gpu)
+
+    convert_weights(model)
+    convert_weights(img2text)
+
+    self.model = model
+    self.img2text = img2text
+
+    # Load the dataset
+    logging.info("Loading dataset.")
+
+    root_project = os.path.join(get_project_root(), _DATA_DIR)
+    dataset = CustomFolder(
+        os.path.join(root_project, _COCO_VAL2017_PATH), transform=preprocess_val
+    )
+
+    # Initialize the dataloader. This is used to create the pickle file from
+    # the dataset.
+    dataloader = DataLoader(
+        dataset,
+        batch_size=64,
+        shuffle=False,
+        num_workers=1,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    self.dataloader = dataloader
+
+    logging.info("Finished initializing Pic2Word server.")
+
+  def preprocess(self, data: Any) -> str:
+    """Preprocess input data."""
+    logging.info("Preprocessing Pic2Word inference request.")
+    query = data[0]
+
+    self.output_storage_dir = query[_REQUEST_OUTPUT_STORAGE_DIR]
+    prompts = query[_REQUEST_PROMPTS]
+    prompts = prompts.split(",")
+    self.prompt = prompts
+
+    image_path = query[_REQUEST_IMAGE_PATH]
+    # The query image is only supported via GCS bucket upload.
+    fileutils.download_gcs_dir_to_local(image_path, _LOCAL_QUERY_PATH)
+    image_file_name = query[_REQUEST_IMAGE_FILE_NAME]
+
+    query_file = f"./query/{image_file_name}"
+
+    logging.info("Setting model args.")
+
+    args = {
+        "openai-pretrained": True,
+        "resume": _CHECKPOINT_DIR,
+        "retrieval_data": _COCO_DATASET_NAME,
+        "query_file": query_file,
+        "demo_out": _OUTPUT_LOCAL_DIR,
+        "prompts": prompts,
+        "distributed": False,
+        "dp": False,
+        "gpu": 0,
+        "model": _MODEL_NAME,
+        "world_size": 1,
+    }
+    model_input = Namespace(**args)
+
+    logging.info("Finished preprocessing Pic2Word inference request.")
+    return model_input
+
+  def inference(self, model_input: Any):
+    """Runs inference."""
+    logging.info("Running model-inference.")
+    visualize_results(
+        model=self.model,
+        img2text=self.img2text,
+        args=model_input,
+        prompt=self.prompt,
+        dataloader=self.dataloader,
+    )
+
+  def postprocess(self):
+    """Upload the output images to the bucket."""
+    logging.info("Running request postprocess.")
+    fileutils.upload_local_dir_to_gcs(
+        _IMAGE_OUTPUT_LOCAL_DIR, self.output_storage_dir
+    )
+
+  def handle(self, data: Any, context: Any) -> str:  # pylint: disable=unused-argument
+    """Runs preprocess, inference, and post-processing."""
+    logging.info("Received Pic2Word inference request")
+    model_input = self.preprocess(data)
+    self.inference(model_input)
+    self.postprocess()
+    logging.info("Done handling input.")
+    return _RESPONSE_MSG

community-content/vertex_vision_model_garden/model_oss/timm/dockerfile/serve.Dockerfile

@@ -0,0 +1,59 @@
+# Dockerfile for serving dockers with timm.
+#
+# To build:
+# docker build -f model_oss/timm/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM pytorch/torchserve:0.7.0-gpu
+
+USER root
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="timm_serving"
+
+# Install timm.
+RUN python3 -m pip install --upgrade pip
+RUN python3 -m pip install timm==0.6.12
+RUN python3 -m pip install google-cloud-storage==2.9.0
+
+# Copy model artifacts.
+COPY model_oss/timm/handler.py /home/model-server/handler.py
+
+# Create torchserve configuration file.
+RUN echo \
+    "default_response_timeout=1200\n" \
+    "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${infer_port}\n"  \
+    "management_address=http://0.0.0.0:${mng_port}" >> /home/model-server/config.properties
+
+# Expose ports.
+EXPOSE ${infer_port}
+EXPOSE ${mng_port}
+
+# Archive eager mode model artifacts and dependencies.
+# Do not set --model-file and --serialized-file because model and checkpoint will be dynamically loaded in handler.py.
+RUN torch-model-archiver \
+--model-name=${model_name} \
+--version=1.0 \
+--handler=/home/model-server/handler.py \
+--runtime=python3 \
+--export-path=/home/model-server/model-store \
+--archive-format=default \
+--force
+
+# Run Torchserve HTTP serve to respond to prediction requests.
+CMD ["torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${model_name}=${model_name}.mar", \
+     "--model-store", "/home/model-server/model-store"]

community-content/vertex_vision_model_garden/model_oss/timm/dockerfile/train.Dockerfile

@@ -0,0 +1,47 @@
+# Dockerfile for basic training dockers with timm.
+#
+# To build:
+# docker build -f model_oss/timm/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+# Base on pytorch-cuda image.
+FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Download timm source code with pinned version.
+RUN wget -q https://github.com/rwightman/pytorch-image-models/archive/refs/tags/v0.6.12.tar.gz
+RUN tar xzf v0.6.12.tar.gz
+
+# Install libraries.
+RUN pip install cloudml-hypertune==0.1.0.dev6
+
+# Switch to timm repo.
+WORKDIR /workspace/pytorch-image-models-0.6.12
+
+# NOTE: use 'sed' to modify the timm source code to
+#       make timm CheckpointSaver can work with gcsfuse.
+RUN sed -i "1 i\import shutil" timm/utils/checkpoint_saver.py
+RUN sed -i "s#os.link#shutil.copyfile#g" timm/utils/checkpoint_saver.py
+RUN sed -i "s#os.unlink#os.remove#g" timm/utils/checkpoint_saver.py
+
+# NOTE: use 'sed' to modify the timm source code to
+#       add hp training support to timm trainer.
+RUN sed -i "693 a\        if saver is not None: hpt = hypertune.HyperTune(); hpt.report_hyperparameter_tuning_metric(hyperparameter_metric_tag='top1_accuracy', metric_value=best_metric, global_step=best_epoch)" train.py
+RUN sed -i "1 i\import hypertune" train.py
+
+# Install timm from source code.
+RUN pip install -e .
+
+# https://pytorch.org/docs/stable/elastic/run.html
+ENTRYPOINT ["torchrun"]

community-content/vertex_vision_model_garden/model_oss/timm/handler.py

@@ -0,0 +1,97 @@
+"""Custom handler for TIMM models."""
+
+import logging
+import os
+from typing import Any
+
+from google.cloud import storage
+import timm
+import torch
+from ts.torch_handler.base_handler import load_label_mapping
+from ts.torch_handler.image_classifier import ImageClassifier
+
+
+GCS_PREFIX = "gs://"
+DOWNLOAD_DIR = "/tmp/download"
+
+
+def download_gcs_file(gcs_uri: str, local_dir: str) -> str:
+  """Download a GCS file to a local directory.
+
+  Arguments:
+    gcs_uri: A string of file path on GCS.
+    local_dir: A string of local directory path.
+
+  Returns:
+    Local path to downloaded file.
+  """
+  if not gcs_uri.startswith(GCS_PREFIX):
+    raise ValueError(f"{gcs_uri} is not a GCS path starting with gs://.")
+
+  file_name = os.path.basename(gcs_uri)
+  local_file_path = os.path.join(local_dir, file_name)
+  os.makedirs(local_dir, exist_ok=True)
+  client = storage.Client()
+  with open(local_file_path, "wb") as f:
+    client.download_blob_to_file(gcs_uri, f)
+  return local_file_path
+
+
+class TimmHandler(ImageClassifier):
+  """Custom handler for TIMM models."""
+
+  def initialize(self, context: Any):
+    """Custom initialize."""
+
+    properties = context.system_properties
+    self.map_location = (
+        "cuda"
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else "cpu"
+    )
+    self.device = torch.device(
+        self.map_location + ":" + str(properties.get("gpu_id"))
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else self.map_location
+    )
+    self.manifest = context.manifest
+
+    # Load timm model by model name.
+    self.model_name = os.environ["MODEL_NAME"]
+    # Whether to use timm pretrained weights, MODEL_PT_PATH overrides this.
+    timm_pretrained = True if os.environ.get("TIMM_PRETRAINED") else False
+    # Load custom checkpoint, it overrides TIMM_PRETRAINED model.
+    self.model_pt_path = os.environ.get("MODEL_PT_PATH")
+    if self.model_pt_path and self.model_pt_path.startswith(GCS_PREFIX):
+      self.model_pt_path = download_gcs_file(self.model_pt_path, DOWNLOAD_DIR)
+
+    if self.model_pt_path and self.model_pt_path.endswith(".pt"):
+      logging.info(
+          "Load model with .pt in jit mode, not working for all timm models"
+          " yet."
+      )
+      self.model = self._load_torchscript_model(self.model_pt_path)
+    else:
+      logging.info("Load model with .pth in eager mode.")
+      self.model = timm.create_model(
+          self.model_name, pretrained=timm_pretrained
+      )
+      if self.model_pt_path and (
+          self.model_pt_path.endswith(".pth")
+          or self.model_pt_path.endswith(".pth.tar")
+      ):
+        checkpoint = torch.load(self.model_pt_path, map_location=self.device)
+        state_dict = checkpoint["state_dict"]
+        self.model.load_state_dict(state_dict)
+      self.model.to(self.device)
+    self.model.eval()
+
+    mapping_file_path = os.environ.get("INDEX_TO_NAME_FILE")
+    if mapping_file_path:
+      if mapping_file_path.startswith(GCS_PREFIX):
+        mapping_file_path = download_gcs_file(mapping_file_path, DOWNLOAD_DIR)
+      self.mapping = load_label_mapping(mapping_file_path)
+
+    self.initialized = True
+
+  # NOTE: Preprocess and postprocess are implemented by ImageClassifier.

community-content/vertex_vision_model_garden/model_oss/util/commons.py

@@ -0,0 +1,79 @@
+"""Common utility lib for prediction on images."""
+
+from typing import Any, Dict, List
+
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+import yaml
+
+from util import image_format_converter
+
+
+def get_prediction_instances(image: Image.Image) -> List[Dict[str, Any]]:
+  """Gets prediction instances.
+
+  Args:
+    image: Image instance.
+
+  Returns:
+    List[Dict[str, Any]]: List of prediction instances.
+  """
+  instances = [{
+      "encoded_image": {"b64": image_format_converter.image_to_base64(image)},
+  }]
+  return instances
+
+
+def get_label_map(label_map_yaml_filepath: str) -> Dict[str, Any]:
+  """Gets the label map from a YAML file.
+
+  Args:
+      label_map_yaml_filepath: Filepath to the label map YAML file.
+
+  Returns:
+      dict: Label map.
+  """
+  with tf.io.gfile.GFile(label_map_yaml_filepath, "rb") as input_file:
+    label_map = yaml.safe_load(input_file.read())
+  return label_map
+
+
+def get_object_detection_endpoint_predictions(
+    detection_endpoint: ...,
+    input_image: np.ndarray,
+    detection_thresh: float = 0.2,
+) -> np.ndarray:
+  """Gets endpoint predictions.
+
+  Args:
+      detection_endpoint: image object detection endpoint.
+      input_image: Input image.
+      detection_thresh: Detection threshold.
+
+  Returns:
+      Object detection predictions from endpoints.
+  """
+  height, width, _ = input_image.shape
+  predictions = detection_endpoint.predict(
+      get_prediction_instances(Image.fromarray(input_image))
+  ).predictions
+  detection_scores = np.array(predictions[0]["detection_scores"])
+  detection_classes = np.array(predictions[0]["detection_classes"])
+  detection_boxes = np.array(
+      [
+          [b[1] * width, b[0] * height, b[3] * width, b[2] * height]
+          for b in predictions[0]["detection_boxes"]
+      ]
+  )
+  thresh_indices = [
+      x for x, val in enumerate(detection_scores) if val > detection_thresh
+  ]
+  preds_merge_conf = np.column_stack((
+      detection_boxes[thresh_indices],
+      detection_scores[thresh_indices],
+  ))
+  preds_merge_cls = np.column_stack(
+      (preds_merge_conf, detection_classes[thresh_indices])
+  )
+  return preds_merge_cls

community-content/vertex_vision_model_garden/model_oss/util/constants.py

@@ -1,9 +1,14 @@
 """Vertex vision model garden util constants."""
 
-# Objectives.
+# TfVision Objectives.
 OBJECTIVE_IMAGE_CLASSIFICATION = 'icn'
 OBJECTIVE_IMAGE_OBJECT_DETECTION = 'iod'
 OBJECTIVE_IMAGE_SEGMENTATION = 'isg'
+OBJECTIVE_VIDEO_CLASSIFICATION = 'vcn'
+OBJECTIVE_VIDEO_ACTION_RECOGNITION = 'var'
+
+# PyTorch Models.
+OBJECTIVE_TIMM = 'timm'
 
 # Input file types.
 INPUT_FILE_TYPE_CSV = 'csv'
@@ -61,4 +66,20 @@ GCSFUSE_URI_PREFIX = '/gcs/'
 
 LOCAL_EVALUATION_RESULT_DIR = '/tmp/evaluation_result_dir'
 LOCAL_MODEL_DIR = '/tmp/model_dir'
+LOCAL_BASE_MODEL_DIR = '/tmp/base_model_dir'
 LOCAL_DATA_DIR = '/tmp/data'
+
+# Huggingface files.
+HF_MODEL_WEIGHTS_SUFFIX = '.bin'
+
+# PEFT finetuning constants.
+TEXT_TO_IMAGE_LORA = 'text-to-image-lora'
+SEQUENCE_CLASSIFICATION_LORA = 'sequence-classification-lora'
+CAUSAL_LANGUAGE_MODELING_LORA = 'causal-language-modeling-lora'
+INSTRUCT_LORA = 'instruct-lora'
+
+# Precision modes for loading model weights.
+PRECISION_MODE_4 = '4bit'
+PRECISION_MODE_8 = '8bit'
+PRECISION_MODE_16 = 'float16'
+PRECISION_MODE_32 = 'float32'

community-content/vertex_vision_model_garden/model_oss/util/fileutils.py

@@ -2,6 +2,10 @@
 
 import glob
 import os
+import pathlib
+import shutil
+from typing import Tuple
+import uuid
 
 from absl import logging
 from google.cloud import storage
@@ -9,6 +13,44 @@ from google.cloud import storage
 from util import constants
 
 
+def generate_tmp_path(extension: str = '') -> str:
+  """Generates a temporary file path with UUID.
+
+  Args:
+    extension: File extension, e.g. '.jpg', '.avi'. If not given, no extension
+      will be appended to the filename.
+
+  Returns:
+    Generated file path.
+  """
+  return os.path.join(constants.LOCAL_DATA_DIR, uuid.uuid1().hex) + extension
+
+
+def force_gcs_fuse_path(gcs_uri: str) -> str:
+  """Converts gs:// uris to their /gcs/ equivalents. No-op for other uris."""
+  if is_gcs_path(gcs_uri):
+    return (
+        constants.GCSFUSE_URI_PREFIX + gcs_uri[len(constants.GCS_URI_PREFIX) :]
+    )
+  else:
+    return gcs_uri
+
+
+def download_gcs_file_to_local_dir(gcs_uri: str, local_dir: str):
+  """Download a gcs file to a local dir.
+
+  Args:
+    gcs_uri: A string of file path on GCS.
+    local_dir: A string of local directory.
+  """
+  if not is_gcs_path(gcs_uri):
+    raise ValueError(
+        f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.'
+    )
+  filename = os.path.basename(gcs_uri)
+  download_gcs_file_to_local(gcs_uri, os.path.join(local_dir, filename))
+
+
 def download_gcs_file_to_local(gcs_uri: str, local_path: str):
   """Download a gcs file to a local path.
 
@@ -16,7 +58,7 @@ def download_gcs_file_to_local(gcs_uri: str, local_path: str):
     gcs_uri: A string of file path on GCS.
     local_path: A string of local file path.
   """
-  if not gcs_uri.startswith(constants.GCS_URI_PREFIX):
+  if not is_gcs_path(gcs_uri):
     raise ValueError(
         f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.'
     )
@@ -26,7 +68,9 @@ def download_gcs_file_to_local(gcs_uri: str, local_path: str):
     client.download_blob_to_file(gcs_uri, f)
 
 
-def download_gcs_dir_to_local(gcs_dir: str, local_dir: str):
+def download_gcs_dir_to_local(
+    gcs_dir: str, local_dir: str, skip_hf_model_bin: bool = False
+):
   """Downloads files in a GCS directory to a local directory.
 
   For example:
@@ -37,7 +81,10 @@ def download_gcs_dir_to_local(gcs_dir: str, local_dir: str):
   Arguments:
     gcs_dir: A string of directory path on GCS.
     local_dir: A string of local directory path.
+    skip_hf_model_bin: True to skip downloading HF model bin files.
   """
+  if not is_gcs_path(gcs_dir):
+    raise ValueError(f'{gcs_dir} is not a GCS path starting with gs://.')
   bucket_name = gcs_dir.split('/')[2]
   prefix = gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/')
   client = storage.Client()
@@ -48,8 +95,16 @@ def download_gcs_dir_to_local(gcs_dir: str, local_dir: str):
     file_path = blob.name[len(prefix) :].strip('/')
     local_file_path = os.path.join(local_dir, file_path)
     os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
-    logging.info('Downloading %s to %s', file_path, local_file_path)
-    blob.download_to_filename(local_file_path)
+    if (
+        file_path.endswith(constants.HF_MODEL_WEIGHTS_SUFFIX)
+        and skip_hf_model_bin
+    ):
+      logging.info('Skip downloading model bin %s', file_path)
+      with open(local_file_path, 'w') as f:
+        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}/{file_path}')
+    else:
+      logging.info('Downloading %s to %s', file_path, local_file_path)
+      blob.download_to_filename(local_file_path)
 
 
 def upload_local_dir_to_gcs(local_dir: str, gcs_dir: str):
@@ -77,3 +132,126 @@ def upload_local_dir_to_gcs(local_dir: str, gcs_dir: str):
       )
       blob = bucket.blob(os.path.join(blob_dir, os.path.basename(local_file)))
       blob.upload_from_filename(local_file)
+
+
+def upload_file_to_gcs_path(
+    source_path: str,
+    destination_uri: str,
+):
+  """Uploads local files to GCS uri.
+
+  After upload the destination_uri will contain the same data as the
+  source_path.
+
+  Args:
+      source_path: Required. Path of the local data to copy to GCS.
+      destination_uri: Required. GCS URI where the data should be uploaded.
+
+  Raises:
+      RuntimeError: When source_path does not exist.
+      GoogleCloudError: When the upload process fails.
+  """
+  source_path_obj = pathlib.Path(source_path)
+  if not source_path_obj.exists():
+    raise RuntimeError(f'Source path does not exist: {source_path}')
+
+  storage_client = storage.Client()
+  source_file_path = source_path
+  destination_file_uri = destination_uri
+  logging.info('Uploading "%s" to "%s"', source_file_path, destination_file_uri)
+  destination_blob = storage.Blob.from_string(
+      destination_file_uri, client=storage_client
+  )
+  destination_blob.upload_from_filename(filename=source_file_path)
+
+
+def is_gcs_path(input_path: str) -> bool:
+  """Checks if the input path is a Google Cloud Storage (GCS) path.
+
+  Args:
+      input_path: The input path to be checked.
+
+  Returns:
+      True if the input path is a GCS path, False otherwise.
+  """
+  return input_path.startswith(constants.GCS_URI_PREFIX)
+
+
+def release_text_assets(
+    output_bucket: str, local_text_file_name: str, remote_text_file_name: str
+) -> None:
+  """Releases text assets.
+
+  Args:
+    output_bucket: gcs output bucket.
+    local_text_file_name: Local text file name.
+    remote_text_file_name: Remote text file name.
+
+  Returns:
+    None
+  """
+  remote_file_path = '{}/{}'.format(output_bucket, remote_text_file_name)
+  logging.info('Uploading "%s" to "%s"', local_text_file_name, remote_file_path)
+  upload_file_to_gcs_path(local_text_file_name, remote_file_path)
+  os.remove(local_text_file_name)
+
+
+def upload_video_from_local_to_gcs(
+    output_bucket: str,
+    local_video_file_name: str,
+    remote_video_file_name: str,
+    temp_local_video_file_name: str,
+) -> None:
+  """Uploads video from local to gcs buckent and releases video assets.
+
+  Args:
+    output_bucket: GCS bucket address.
+    local_video_file_name: Local video file name.
+    remote_video_file_name: Remote video file name.
+    temp_local_video_file_name: Temporary local video file name.
+
+  Returns:
+    None
+  """
+  upload_file_to_gcs_path(
+      temp_local_video_file_name,
+      '{}/{}'.format(output_bucket, remote_video_file_name),
+  )
+  shutil.rmtree(local_video_file_name, ignore_errors=True)
+  shutil.rmtree(temp_local_video_file_name, ignore_errors=True)
+
+
+def download_video_from_gcs_to_local(video_file_path: str) -> Tuple[str, str]:
+  """Downloads video from gcs to local folders.
+
+  Args:
+    video_file_path: Path to the video file.
+
+  Returns:
+    Local and remote video file paths.
+  """
+  _, local_video_file_name = os.path.split(video_file_path)
+  file_extension = os.path.splitext(video_file_path)[1]
+  remote_video_file_name = local_video_file_name.replace(
+      file_extension, '_overlay.mp4'
+  )
+  local_file_path = generate_tmp_path(os.path.splitext(video_file_path)[1])
+  logging.info('Downloading %s to %s...', video_file_path, local_file_path)
+  download_gcs_file_to_local(video_file_path, local_file_path)
+  return local_file_path, remote_video_file_name
+
+
+def get_output_video_file(video_output_file_path: str) -> str:
+  """Gets the output video file name for writing video.
+
+  Args:
+    video_output_file_path: Path to the video output file.
+
+  Returns:
+    str: Local video output file path.
+  """
+  file_extension = os.path.splitext(video_output_file_path)[1]
+  out_local_video_file_name = video_output_file_path.replace(
+      file_extension, '_overlay' + file_extension
+  )
+  return out_local_video_file_name

community-content/vertex_vision_model_garden/model_oss/util/hypertune_utils.py

@@ -4,6 +4,7 @@ import os
 
 from absl import logging
 
+
 _ENVIRONMENT_VARIABLE_FOR_TRIAL_ID = 'CLOUD_ML_TRIAL_ID'
 
 

community-content/vertex_vision_model_garden/model_oss/video_util/video_format_converter.py

@@ -1,14 +0,0 @@
-"""Video format converter util lib."""
-
-import io
-from typing import Sequence
-import imageio
-import numpy as np
-from PIL import Image
-
-
-def frames_to_video_bytes(frames: Sequence[np.ndarray], fps: int) -> bytes:
-  images = [Image.fromarray(array) for array in frames]
-  io_obj = io.BytesIO()
-  imageio.mimsave(io_obj, images, format=".mp4", fps=fps)
-  return io_obj.getvalue()

community-content/vertex_vision_model_garden/model_oss/vllm/dockerfile/serve.Dockerfile

@@ -0,0 +1,73 @@
+# Dockerfile for vLLM serving.
+#
+# To build:
+# docker build -f model_oss/vllm/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/{YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/{YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+# The base image is required by vllm
+# https://vllm.readthedocs.io/en/latest/getting_started/installation.html
+FROM nvcr.io/nvidia/pytorch:22.12-py3
+
+USER root
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+RUN apt-get install -y --no-install-recommends apt-utils
+RUN apt-get install -y --no-install-recommends curl
+RUN apt-get install -y --no-install-recommends wget
+RUN apt-get install -y --no-install-recommends git
+RUN apt-get install -y --no-install-recommends jq
+RUN apt-get install -y --no-install-recommends gnupg
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
+
+# Install pytorch
+RUN pip install --upgrade torch==2.0.1
+
+# Install vllm deps.
+RUN pip install xformers==0.0.20
+RUN pip install ninja==1.11.1
+RUN pip install psutil==5.9.5
+RUN pip install ray==2.6.2
+RUN pip install sentencepiece==0.1.99
+RUN pip install fastapi==0.100.1
+RUN pip install uvicorn==0.23.2
+RUN pip install pydantic==1.10.12
+
+# Install transformers from source.
+WORKDIR /workspace
+RUN git clone https://github.com/huggingface/transformers.git
+WORKDIR transformers
+# Pin the commit to add-code-llama at 08/25/2023
+RUN git reset --hard 015f8e110d270a0ad42de4ae5b98198d69eb1964
+RUN pip install -e .
+WORKDIR /workspace
+
+# Install vllm from source.
+RUN git clone https://github.com/vllm-project/vllm.git
+WORKDIR vllm
+# Pin the version to a fixed git commit on 08/16/2023.
+RUN git reset --hard d1744376ae9fdbfa6a2dc763e1c67309e138fa3d
+# Apply a patch to vllm source:
+# 1) For models on Huggingface hub: if the model has multiple bin files, each
+#    bin file is downloaded separately and gets deleted after loading to GPU
+# 2) For models on GCS bucket: each model bin files is download separately
+#    and gets deleted after loading to GPU.
+# 3) Support code-llama model loading.
+COPY model_oss/vllm/vllm.patch /tmp/vllm.patch
+RUN git apply /tmp/vllm.patch
+RUN pip install -e .
+
+# Expose port 7080 for host serving.
+EXPOSE 7080

community-content/vertex_vision_model_garden/model_oss/vllm/vllm.patch

@@ -0,0 +1,311 @@
+diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
+index 99fe593..e11246b 100644
+--- a/vllm/engine/arg_utils.py
++++ b/vllm/engine/arg_utils.py
+@@ -1,12 +1,43 @@
+ import argparse
+ import dataclasses
+ from dataclasses import dataclass
++import os
+ from typing import Optional, Tuple
+ 
++from google.cloud import storage
+ from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
+                          SchedulerConfig)
+ 
+ 
++GCS_PREFIX = "gs://"
++
++
++def is_gcs_path(input_path: str) -> bool:
++    return input_path.startswith(GCS_PREFIX)
++
++
++def download_gcs_dir_to_local(gcs_dir: str, local_dir: str):
++    if os.path.isdir(local_dir):
++        return
++    # gs://bucket_name/dir
++    bucket_name = gcs_dir.split('/')[2]
++    prefix = gcs_dir[len(GCS_PREFIX + bucket_name) :].strip('/')
++    client = storage.Client()
++    blobs = client.list_blobs(bucket_name, prefix=prefix)
++    for blob in blobs:
++        if blob.name[-1] == '/':
++            continue
++        file_path = blob.name[len(prefix) :].strip('/')
++        local_file_path = os.path.join(local_dir, file_path)
++        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
++        if file_path.endswith(".bin"):
++            with open(local_file_path, 'w') as f:
++                f.write(f'{GCS_PREFIX}{bucket_name}/{prefix}/{file_path}')
++        else:
++            print(f"==> Download {gcs_dir}/{file_path} to {local_file_path}")
++            blob.download_to_filename(local_file_path)
++
++
+ @dataclass
+ class EngineArgs:
+     """Arguments for vLLM engine."""
+@@ -143,6 +174,19 @@ class EngineArgs:
+     def create_engine_configs(
+         self,
+     ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]:
++        # Preprocess GCS paths.
++        if is_gcs_path(self.tokenizer) and self.tokenizer != self.model:
++            local_dir = "/tmp/gcs_tokenizer"
++            download_gcs_dir_to_local(self.tokenizer, local_dir)
++            self.tokenizer = local_dir
++        if is_gcs_path(self.model):
++            # Download GCS model without bin files.
++            local_dir = "/tmp/gcs_model"
++            download_gcs_dir_to_local(self.model, local_dir)
++            if self.tokenizer == self.model:
++                self.tokenizer = local_dir
++            self.model = local_dir
++
+         # Initialize the configs.
+         model_config = ModelConfig(self.model, self.tokenizer,
+                                    self.tokenizer_mode, self.trust_remote_code,
+diff --git a/vllm/entrypoints/api_server.py b/vllm/entrypoints/api_server.py
+index 58ea2e2..350e209 100644
+--- a/vllm/entrypoints/api_server.py
++++ b/vllm/entrypoints/api_server.py
+@@ -15,6 +15,10 @@ TIMEOUT_KEEP_ALIVE = 5  # seconds.
+ TIMEOUT_TO_PREVENT_DEADLOCK = 1  # seconds.
+ app = FastAPI()
+ 
++# Required by Vertex deployment.
++@app.get("/ping")
++async def ping() -> Response:
++    return Response(status_code=200)
+ 
+ @app.post("/generate")
+ async def generate(request: Request) -> Response:
+@@ -26,6 +30,9 @@ async def generate(request: Request) -> Response:
+     - other fields: the sampling parameters (See `SamplingParams` for details).
+     """
+     request_dict = await request.json()
++    is_on_vertex = "instances" in request_dict
++    if is_on_vertex:
++        request_dict = request_dict["instances"][0]
+     prompt = request_dict.pop("prompt")
+     stream = request_dict.pop("stream", False)
+     sampling_params = SamplingParams(**request_dict)
+@@ -63,7 +70,10 @@ async def generate(request: Request) -> Response:
+     assert final_output is not None
+     prompt = final_output.prompt
+     text_outputs = [prompt + output.text for output in final_output.outputs]
+-    ret = {"text": text_outputs}
++    if is_on_vertex:
++        ret = {"predictions": text_outputs}
++    else:
++        ret = {"text": text_outputs}
+     return JSONResponse(ret)
+ 
+ 
+diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py
+index 93ab499..eca1b89 100644
+--- a/vllm/model_executor/models/llama.py
++++ b/vllm/model_executor/models/llama.py
+@@ -85,6 +85,7 @@ class LlamaAttention(nn.Module):
+         hidden_size: int,
+         num_heads: int,
+         num_kv_heads: int,
++        rope_theta: float = 10000,
+     ):
+         super().__init__()
+         self.hidden_size = hidden_size
+@@ -99,6 +100,7 @@ class LlamaAttention(nn.Module):
+         self.q_size = self.num_heads * self.head_dim
+         self.kv_size = self.num_kv_heads * self.head_dim
+         self.scaling = self.head_dim**-0.5
++        self.rope_theta = rope_theta
+ 
+         self.qkv_proj = ColumnParallelLinear(
+             hidden_size,
+@@ -118,6 +120,7 @@ class LlamaAttention(nn.Module):
+         self.attn = PagedAttentionWithRoPE(self.num_heads,
+                                            self.head_dim,
+                                            self.scaling,
++                                           base=self.rope_theta,
+                                            rotary_dim=self.head_dim,
+                                            num_kv_heads=self.num_kv_heads)
+ 
+@@ -143,10 +146,15 @@ class LlamaDecoderLayer(nn.Module):
+     def __init__(self, config: LlamaConfig):
+         super().__init__()
+         self.hidden_size = config.hidden_size
++        try:
++            rope_theta = config.rope_theta
++        except AttributeError:
++            rope_theta = 10000
+         self.self_attn = LlamaAttention(
+             hidden_size=self.hidden_size,
+             num_heads=config.num_attention_heads,
+             num_kv_heads=config.num_key_value_heads,
++            rope_theta=rope_theta,
+         )
+         self.mlp = LlamaMLP(
+             hidden_size=self.hidden_size,
+diff --git a/vllm/model_executor/weight_utils.py b/vllm/model_executor/weight_utils.py
+index a9d899a..57f39b5 100644
+--- a/vllm/model_executor/weight_utils.py
++++ b/vllm/model_executor/weight_utils.py
+@@ -3,13 +3,17 @@ import filelock
+ import glob
+ import json
+ import os
++import time
+ from typing import Iterator, List, Optional, Tuple
+ 
+-from huggingface_hub import snapshot_download
++from google.cloud import storage
++from huggingface_hub import hf_hub_download, snapshot_download
+ import numpy as np
+ import torch
+ from tqdm.auto import tqdm
+ 
++HF_PREFIX = "hf://"
++
+ 
+ class Disabledtqdm(tqdm):
+ 
+@@ -22,60 +26,90 @@ def hf_model_weights_iterator(
+     cache_dir: Optional[str] = None,
+     use_np_cache: bool = False,
+ ) -> Iterator[Tuple[str, torch.Tensor]]:
++    if use_np_cache:
++        raise ValueError("Do not support use_np_cache for lazy download.")
++
+     # Prepare file lock directory to prevent multiple processes from
+     # downloading the same model weights at the same time.
+     lock_dir = cache_dir if cache_dir is not None else "/tmp"
+     lock_file_name = model_name_or_path.replace("/", "-") + ".lock"
+     lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name))
+ 
+-    # Download model weights from huggingface.
+-    is_local = os.path.isdir(model_name_or_path)
+-    if not is_local:
+-        with lock:
+-            hf_folder = snapshot_download(model_name_or_path,
+-                                          allow_patterns="*.bin",
+-                                          cache_dir=cache_dir,
+-                                          tqdm_class=Disabledtqdm)
+-    else:
+-        hf_folder = model_name_or_path
+-
+-    hf_bin_files = [
+-        x for x in glob.glob(os.path.join(hf_folder, "*.bin"))
+-        if not x.endswith("training_args.bin")
+-    ]
+-
+-    if use_np_cache:
+-        # Convert the model weights from torch tensors to numpy arrays for
+-        # faster loading.
+-        np_folder = os.path.join(hf_folder, "np")
+-        os.makedirs(np_folder, exist_ok=True)
+-        weight_names_file = os.path.join(np_folder, "weight_names.json")
+-        with lock:
+-            if not os.path.exists(weight_names_file):
+-                weight_names = []
+-                for bin_file in hf_bin_files:
+-                    state = torch.load(bin_file, map_location="cpu")
+-                    for name, param in state.items():
+-                        param_path = os.path.join(np_folder, name)
+-                        with open(param_path, "wb") as f:
+-                            np.save(f, param.cpu().detach().numpy())
+-                        weight_names.append(name)
+-                with open(weight_names_file, "w") as f:
+-                    json.dump(weight_names, f)
+-
+-        with open(weight_names_file, "r") as f:
+-            weight_names = json.load(f)
+-
+-        for name in weight_names:
+-            param_path = os.path.join(np_folder, name)
+-            with open(param_path, "rb") as f:
+-                param = np.load(f)
+-            yield name, torch.from_numpy(param)
++    bin_files = []
++    if not os.path.isdir(model_name_or_path):
++        try:
++            with lock:
++                index_file = hf_hub_download(repo_id=model_name_or_path,
++                                             filename="pytorch_model.bin.index.json",
++                                             cache_dir=cache_dir)
++        except:
++            print("==> The model is in HF hub with 1 bin file, download it directly.", flush=True)
++            with lock:
++                hf_folder = snapshot_download(repo_id=model_name_or_path,
++                                              allow_patterns="*.bin",
++                                              cache_dir=cache_dir,
++                                              tqdm_class=Disabledtqdm)
++            bin_files = [x for x in glob.glob(os.path.join(hf_folder, "*.bin"))]
++        else:
++            print("==> The model is in HF hub with multiple bin file, do not download it now.", flush=True)
++            with open(index_file, "r") as f:
++                index = json.loads(f.read())
++            bin_filenames = set(index["weight_map"].values())
++            bin_files = [f"{HF_PREFIX}{model_name_or_path}/{bin_filename}" for bin_filename in bin_filenames]
+     else:
+-        for bin_file in hf_bin_files:
+-            state = torch.load(bin_file, map_location="cpu")
+-            for name, param in state.items():
+-                yield name, param
++        print("==> The model is in local disk.", flush=True)
++        bin_files = [x for x in glob.glob(os.path.join(model_name_or_path, "*.bin"))]
++
++    if "training_args.bin" in bin_files:
++        bin_files.remove("training_args.bin")
++    bin_files.sort()
++    print(f"==> Fetched bin files: {bin_files}", flush=True)
++
++    model_dir = "/tmp/model"
++    os.makedirs(model_dir, exist_ok=True)
++    for bin_file in bin_files:
++        delete_download = False
++
++        if os.path.exists(bin_file):
++            if open(bin_file, "rb").read(2) == b"gs":
++                gcs_path = open(bin_file).read()
++                bin_filename = gcs_path.split("/")[-1]
++                local_file = os.path.join(model_dir, bin_filename)
++                with lock:
++                    if not os.path.exists(local_file):
++                        client = storage.Client()
++                        with open(local_file, 'wb') as f:
++                            print(f"==> Download {gcs_path} to {bin_file}", flush=True)
++                            client.download_blob_to_file(gcs_path, f)
++                bin_file = local_file
++                delete_download = True
++        else:
++            assert bin_file.startswith(HF_PREFIX)
++            bin_filename = os.path.basename(bin_file)
++            local_file = os.path.join(model_dir, bin_filename)
++            with lock:
++                if not os.path.exists(local_file):
++                    print(f"==> Download {model_name_or_path}/{bin_filename} to {local_file}", flush=True)
++                    hf_hub_download(repo_id=model_name_or_path,
++                                    filename=bin_filename,
++                                    local_dir=model_dir,
++                                    local_dir_use_symlinks=False,
++                                    force_download=True)
++            bin_file = local_file
++            delete_download = True
++
++        torch.distributed.barrier()
++        print(f"==> Load {bin_file} to memory.", flush=True)
++        state = torch.load(bin_file, map_location="cpu")
++        for name, param in state.items():
++            yield name, param
++        torch.distributed.barrier()
++
++        if delete_download:
++            with lock:
++                if os.path.exists(bin_file):
++                    print(f"==> Delete {bin_file}", flush=True)
++                    os.remove(bin_file)
+ 
+ 
+ def load_tensor_parallel_weights(
+ 

community-content/vertex_vision_model_garden/model_oss/vot/dockerfile/serve.Dockerfile

@@ -0,0 +1,87 @@
+# Dockerfile for basic serving dockers for vot.
+#
+# To build:
+# docker build -f model_oss/vot/dockerfile/serving.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+# Switch to this base image for gpu serve.
+FROM pytorch/torchserve:0.7.0-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="vot_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/vot"
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install accelerate==0.17.0
+RUN pip install datasets==2.9.0
+RUN pip install bytetracker==0.3.2
+RUN pip install imageio[ffmpeg]==2.31.1
+RUN pip install google-cloud-aiplatform==1.25.0
+RUN pip install google-cloud-storage==2.9.0
+RUN pip install fastapi==0.96.0
+RUN pip install lap==0.4.0
+RUN pip install numpy==1.24.3
+RUN pip install opencv-python==4.7.0.72
+RUN pip install Pillow==9.5.0
+RUN pip install protobuf==3.19.6
+RUN pip install pandas==2.0.2
+RUN pip install pycocotools==2.0.6
+RUN pip install scipy==1.10.1
+RUN pip install tensorflow==2.11.1
+RUN pip install torch==2.0.1
+RUN pip install torchvision==0.15.2
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install uvicorn==0.22.0
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY model_oss/vot/handler.py /home/model-server/handler.py
+COPY model_oss/vot/visualization_utils.py /home/model-server/vot/
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server/
+
+# Create torchserve configuration file.
+RUN echo \
+    "default_response_timeout=3600\n" \
+    "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${infer_port}\n"  \
+    "management_address=http://0.0.0.0:${mng_port}" >> /home/model-server/config.properties
+
+# Expose ports.
+EXPOSE ${infer_port}
+EXPOSE ${mng_port}
+
+# Archive model artifacts and dependencies.
+# Do not set --model-file and --serialized-file because model and checkpoint will be dynamically loaded in handler.py.
+RUN torch-model-archiver \
+--model-name=${model_name} \
+--version=1.0 \
+--handler=/home/model-server/handler.py \
+--runtime=python3 \
+--export-path=/home/model-server/model-store \
+--archive-format=default \
+--force
+
+# Run Torchserve HTTP serve to respond to prediction requests.
+CMD ["torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${model_name}=${model_name}.mar", \
+     "--model-store", "/home/model-server/model-store"]

community-content/vertex_vision_model_garden/model_oss/vot/handler.py

@@ -0,0 +1,196 @@
+"""Custom handler for video object tracking models."""
+
+import logging
+import os
+import tempfile
+from typing import Any, List, Optional, Tuple
+
+from bytetracker import BYTETracker
+import cv2
+from google.cloud import aiplatform
+import imageio.v2 as iio
+from PIL import Image
+import tensorflow as tf
+import torch
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import commons
+from util import fileutils
+import visualization_utils
+
+_VIDEO_URI = "video_uri"
+_DATA = "data"
+_TRACK_THRESHOLD = 0.45
+_TRACK_BUFFER = 25
+_MATCH_THRESHOLD = 0.8
+
+
+class VideoObjectTrackingHandler(BaseHandler):
+  """Custom handler for video object tracking models."""
+
+  def initialize(self, context: Any) -> None:
+    properties = context.system_properties
+    self.map_location = (
+        "cuda"
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else "cpu"
+    )
+    self.device = torch.device(
+        self.map_location + ":" + str(properties.get("gpu_id"))
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else self.map_location
+    )
+    self.manifest = context.manifest
+
+    detection_endpoint_id = os.environ.get("DETECTION_ENDPOINT", None)
+    if detection_endpoint_id:
+      self.detection_endpoint = aiplatform.Endpoint(detection_endpoint_id)
+      endpoint_label_map = os.environ.get("LABEL_MAP", None)
+      if endpoint_label_map:
+        endpoint_label_map_file = endpoint_label_map
+        self.label_map = commons.get_label_map(endpoint_label_map_file)
+      else:
+        raise ValueError(
+            "LABEL MAP must be provided with DETECTION ENDPOINT:"
+            f" {self.detection_endpoint}"
+        )
+
+    self.track_thresh = os.environ.get("TRACK_THRESHOLD", _TRACK_THRESHOLD)
+    self.track_buffer = os.environ.get("TRACK_BUFFER", _TRACK_BUFFER)
+    self.match_thresh = os.environ.get("MATCH_THRESHOLD", _MATCH_THRESHOLD)
+    self.save_video_results = bool(int(os.environ.get("SAVE_VIDEO_RESULTS", 0)))
+    self.output_bucket = os.environ.get("OUTPUT_BUCKET", None)
+    if not self.output_bucket:
+      raise ValueError("Empty Output Bucket.")
+    self.initialized = True
+    logging.info("Handler initialization done.")
+
+  def preprocess(
+      self, data: Any
+  ) -> Tuple[Optional[List[str]], Optional[List[Image.Image]]]:
+    """Preprocesses the input data.
+
+    Args:
+        data (Any): Input data.
+
+    Returns:
+        List of videos uris.
+    """
+    video_uris = None
+    if _VIDEO_URI in data[0]:
+      video_uris = [item[_VIDEO_URI] for item in data]
+    # TorchServe's default handlers expect each instance
+    # to be wrapped in a data field for batch prediction.
+    if _DATA in data[0]:
+      video_uris = [item[_DATA][_VIDEO_URI] for item in data]
+    return video_uris
+
+  def inference(self, data: Any, *args, **kwargs) -> List[Any]:
+    """Runs object detection and tracking inference on a video frame by frame.
+
+    If using yolo detection, the function uses the ultralytics yolo models for
+    IOD, otherwise it uses the provided IOD endpoint and associated the selected
+    tracking method to the detections.
+
+    Args:
+        data: List of video files.
+        *args: Additional arguments.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        List of video frame annotations and/or output decorated video uris.
+    """
+    gcs_video_files = data
+    video_preds = []
+    for gcs_video_file in gcs_video_files:
+      results_info = {}
+      temp_text_file = tempfile.NamedTemporaryFile(delete=False, mode="w+t")
+      local_video_file_name, remote_video_file_name = (
+          fileutils.download_video_from_gcs_to_local(gcs_video_file)
+      )
+      remote_text_file_name = remote_video_file_name.replace(
+          "overlay.mp4", "annotations.txt"
+      )
+
+      cap = cv2.VideoCapture(local_video_file_name)
+      fps = cap.get(cv2.CAP_PROP_FPS)
+      temp_local_video_file_name = fileutils.get_output_video_file(
+          local_video_file_name
+      )
+      if self.save_video_results:
+        self.video_writer = iio.get_writer(
+            temp_local_video_file_name,
+            format="FFMPEG",
+            mode="I",
+            fps=float(fps),
+            codec="h264",
+        )
+      self.tracker = BYTETracker(
+          track_thresh=self.track_thresh,
+          track_buffer=self.track_buffer,
+          match_thresh=self.match_thresh,
+          frame_rate=fps,
+      )
+      frame_idx = 1
+      while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+          break
+
+        dets_np = commons.get_object_detection_endpoint_predictions(
+            self.detection_endpoint, frame
+        )
+        dets_tf = tf.convert_to_tensor(dets_np)
+        online_targets = self.tracker.update(dets_tf, None)
+        if online_targets.size > 0:
+          frame = visualization_utils.overlay_tracking_results(
+              frame_idx,
+              frame,
+              online_targets,
+              label_map=self.label_map,
+              temp_text_file_path=temp_text_file.name,
+          )
+
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        if self.save_video_results:
+          self.video_writer.append_data(frame)
+        logging.info(
+            "Finished processing frame %s for video %s.",
+            frame_idx,
+            gcs_video_file,
+        )
+        frame_idx += 1
+
+      self.video_writer.close()
+      cap.release()
+      if self.save_video_results:
+        fileutils.upload_video_from_local_to_gcs(
+            self.output_bucket,
+            local_video_file_name,
+            remote_video_file_name,
+            temp_local_video_file_name,
+        )
+        results_info["output_video"] = "{}/{}".format(
+            self.output_bucket, remote_video_file_name
+        )
+
+      fileutils.release_text_assets(
+          self.output_bucket,
+          temp_text_file.name,
+          remote_text_file_name,
+      )
+      results_info["annotations"] = "{}/{}".format(
+          self.output_bucket, remote_text_file_name
+      )
+      video_preds.append(results_info)
+
+    return video_preds
+
+  def handle(self, data: Any, context: Any) -> List[Any]:
+    model_input = self.preprocess(data)
+    model_out = self.inference(model_input)
+    output = self.postprocess(model_out)
+    return output
+
+  def postprocess(self, inference_result: List[Any]) -> List[Any]:
+    return inference_result

community-content/vertex_vision_model_garden/model_oss/vot/visualization_utils.py

@@ -0,0 +1,168 @@
+"""Image and bounding box visualization util lib."""
+
+from typing import Dict, List, Optional
+import cv2
+import numpy as np
+from PIL import ImageColor
+
+
+def draw_bounding_box_on_image(
+    image: np.ndarray,
+    ymin: float,
+    xmin: float,
+    ymax: float,
+    xmax: float,
+    color: str,
+    thickness: int = 4,
+    display_str_list: Optional[List[str]] = None,
+) -> np.ndarray:
+  """Draws a bounding box on an image.
+
+  Args:
+    image: The image to draw the bounding box on.
+    ymin: The minimum y-coordinate of the bounding box.
+    xmin: The minimum x-coordinate of the bounding box.
+    ymax: The maximum y-coordinate of the bounding box.
+    xmax: The maximum x-coordinate of the bounding box.
+    color: The color of the bounding box.
+    thickness: The thickness of the bounding box lines. Defaults to 4.
+    display_str_list: List of strings to display in new line inside the bounding
+      box.
+
+  Returns:
+    An image with a bounding box.
+  """
+  color = ImageColor.getrgb(color)
+  cv2.rectangle(
+      image, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, thickness
+  )
+  # Display the strings below the bounding box
+  for i, display_str in enumerate(display_str_list):
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    scale = 0.4
+    thickness = 1
+    text_width, text_height = cv2.getTextSize(
+        display_str, font, scale, thickness
+    )[0]
+    text_bottom = int(ymin - i * text_height)
+    text_left = int(xmin)
+    cv2.rectangle(
+        image,
+        (text_left, text_bottom - text_height),
+        (text_left + text_width, text_bottom),
+        color,
+        -1,
+    )
+    cv2.putText(
+        image,
+        display_str,
+        (text_left, text_bottom),
+        font,
+        scale,
+        (0, 0, 0),
+        thickness,
+    )
+
+  return image
+
+
+def draw_boxes(
+    image: np.ndarray,
+    boxes: List[List[float]],
+    track_ids: List[int],
+    class_names: List[str],
+    scores: List[float],
+    max_boxes: int = 40,
+    min_score: float = 0.05,
+) -> np.ndarray:
+  """Overlays labeled boxes on an image with formatted scores and label names.
+
+  Args:
+      image: The image to overlay the boxes on.
+      boxes: List of bounding box coordinates [xmin, ymin, xmax, ymax].
+      track_ids: List of track IDs corresponding to each box.
+      class_names: List of class names corresponding to each box.
+      scores: List of scores corresponding to each box.
+      max_boxes: Maximum number of boxes to draw. Defaults to 40.
+      min_score: Minimum score threshold for displaying a box. Defaults to 0.05.
+
+  Returns:
+      PIL.Image.Image: The image with the labeled boxes overlay.
+  """
+  colors = list(ImageColor.colormap.values())
+  for i in range(min(len(boxes), max_boxes)):
+    if scores[i] >= min_score:
+      xmin, ymin, xmax, ymax = boxes[i]
+      display_str = "{}-{}: {}%".format(
+          track_ids[i], class_names[i], int(100 * scores[i])
+      )
+      color = colors[hash(class_names[i]) % len(colors)]
+      image = draw_bounding_box_on_image(
+          image,
+          ymin,
+          xmin,
+          ymax,
+          xmax,
+          color,
+          display_str_list=[display_str],
+      )
+  return image
+
+
+def overlay_tracking_results(
+    frame_idx: int,
+    image_np: np.ndarray,
+    tracker_outputs: np.ndarray,
+    model_names: Optional[Dict[int, str]] = None,
+    label_map: Optional[Dict[str, Dict[int, str]]] = None,
+    temp_text_file_path: Optional[str] = None,
+) -> np.ndarray:
+  """Overlays the results on the image.
+
+  Args:
+      frame_idx: frame index.
+      image_np: Input image.
+      tracker_outputs: Tracker outputs.
+      model_names: label map for yolo models.
+      label_map: label map for IOD detector model.
+      temp_text_file_path: tempfile to save annotations.
+
+  Returns:
+      Decorated output frame.
+  """
+  dboxes = tracker_outputs[:, :4]
+  dtracks = tracker_outputs[:, 4]
+  dclasses = tracker_outputs[:, 5]
+  dscores = tracker_outputs[:, 6]
+  dclasses_as_text = []
+  for detection_class in dclasses:
+    if model_names:
+      dclasses_as_text.append(model_names[int(detection_class)])
+    elif label_map:
+      dclasses_as_text.append(label_map["label_map"][int(detection_class)])
+    else:
+      dclasses_as_text.append("")
+
+  plotted_img = np.array(
+      draw_boxes(
+          image=image_np,
+          boxes=dboxes,
+          track_ids=dtracks,
+          class_names=dclasses_as_text,
+          scores=dscores,
+      )
+  )
+  track_anno_list = []
+  for i, box in enumerate(dboxes):
+    result_list = [dtracks[i], dscores[i], dclasses[i]]
+    xyxy_anno = [np.round(item.item(), 2) for item in box]
+    tracks_anno = (
+        [frame_idx]
+        + [np.round(item.item(), 2) for item in result_list]
+        + xyxy_anno
+    )
+    track_anno_list.append(tracks_anno)
+    with open(temp_text_file_path, "a") as file:
+      file.write(", ".join([str(item) for item in tracks_anno]) + "\n")
+
+  return plotted_img

notebooks/community/CODEOWNERS

@@ -43,8 +43,12 @@
 /notebooks/community/pipelines/google_cloud_pipeline_components_ready_to_go_text_classification_pipeline.ipynb @Narwhalprime
 /notebooks/community/feature_store/get_started_vertex_feature_store.ipynb @junkourata
 /notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb @dstnluong-google
+/notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.pynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_proprietary_image_classification.ipynb @weigary
 /notebooks/community/model_garden/model_garden_proprietary_image_object_detection.ipynb @weigary
 /notebooks/community/model_garden/model_garden_tfvision_image_classification.ipynb @genquan9
@@ -53,6 +57,7 @@
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_2_1.ipynb @bingatgoogle
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_inpainting.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_1_0.ipynb @bingatgoogle
 /notebooks/community/model_garden/model_garden_pytorch_instructpix2pix.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_controlnet.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb @xiangxu-google
@@ -66,11 +71,23 @@
 /notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_pytorch_bart_large_cnn.ipynb @lavraicse
+/notebooks/community/model_garden/model_garden_pytorch_starcoder.ipynb @xcchen1
 /notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb @lavraicse
+/notebooks/community/model_garden/model_garden_jax_fvlm.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_pytorch_text_to_video_zero_shot.ipynb @bingatgoogle
 /notebooks/community/model_garden/model_garden_pytorch_text_to_video.ipynb @KCFindstr
+/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb @xqr-g
 /notebooks/community/generative_ai/text_embedding_api_semantic_search_with_scann.ipynb @henrytansetiawan
 /notebooks/community/bigquery_ml_inference/bq_ml_with_vision_translation_nlp.ipynb @deaconsmith
 /notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_keras_yolov8.ipynb @@dstnluong-google
 /notebooks/community/model_garden/model_garden_pytorch_sam.ipynb @huguensjean
+/notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb @jismailyan
 /notebooks/community/model_garden/model_garden_pytorch_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb @lydhr
+/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb @xiangxu-google

notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb

@@ -0,0 +1,363 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/xqr-g/vertex-ai-samples/blob/main/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "KSP1duKDeaDR"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "JAPoU8Sm5E6e"
+      },
+      "source": [
+        "# Cloud Next Embedding models\n",
+        "\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>                                                                                               \n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "24743cf4a1e1"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.10"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This colab is used as a code example for how to call our newly released text embedding models (textembedding-gecko@latest and textembedding-gecko-multilingual@latest).\n",
+        "\n",
+        "Learn more about [text embedding api](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings).\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "- Vertex LLM SDK\n",
+        "\n",
+        "The steps performed include:\n",
+        "- Installation and imports\n",
+        "- Generate embeddings\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aed92deeb4a0"
+      },
+      "source": [
+        "### 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),\n",
+        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+        "\n",
+        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sBCra4QMA2wR"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FyyMdUeAJIVv"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook.\n",
+        "\n",
+        "**Remember to restart the runtime after installation.**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "snBUuUamoJPz"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install git+https://github.com/googleapis/python-aiplatform.git"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WX3CHZitmSJM"
+      },
+      "source": [
+        "### Please restart the runtime."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dae340cb-0583-4e7e-a562-6817ee4d7f6d"
+      },
+      "source": [
+        "### Imports libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "412d00f1-08db-4880-8ced-52a9583757b8"
+      },
+      "outputs": [],
+      "source": [
+        "import vertexai\n",
+        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "MyMXIZoRlUcR"
+      },
+      "source": [
+        "#### Set your project ID and initiate Vertex AI\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3EdtdqnoldX4"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "REGION = \"us-central1\"\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}\n",
+        "\n",
+        "# Initiate Vertex AI\n",
+        "vertexai.init(project=PROJECT_ID, location=REGION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b"
+      },
+      "source": [
+        "## Generate embeddings"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "hQZoBXNGjizH"
+      },
+      "outputs": [],
+      "source": [
+        "# Set the model name.\n",
+        "MODEL_NAME = \"textembedding-gecko@latest\"  # @param [\"textembedding-gecko@latest\", \"textembedding-gecko-multilingual@latest\"]\n",
+        "\n",
+        "# Set the task_type, text and optional title as the model inputs.\n",
+        "TASK_TYPE = \"RETRIEVAL_DOCUMENT\"  # @param [\"RETRIEVAL_QUERY\", \"RETRIEVAL_DOCUMENT\", \"SEMANTIC_SIMILARITY\", \"CLASSIFICATION\", \"CLUSTERING\"]\n",
+        "TITLE = \"Google\"  # @param {type:\"string\"}\n",
+        "TEXT = \"Embed text.\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Verify the input is valid.\n",
+        "if not MODEL_NAME:\n",
+        "    raise ValueError(\"Please set MODEL_NAME.\")\n",
+        "if not TASK_TYPE:\n",
+        "    raise ValueError(\"Please set TASK_TYPE.\")\n",
+        "if not TEXT:\n",
+        "    raise ValueError(\"Please set TEXT.\")\n",
+        "if TITLE and TASK_TYPE != \"RETRIEVAL_DOCUMENT\":\n",
+        "    raise ValueError(\"Title can only be provided if the task_type is RETRIEVAL_DOCUMENT\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BNPapKXviHlE"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "768\n"
+          ]
+        }
+      ],
+      "source": [
+        "def text_embedding(\n",
+        "  model_name: str, task_type: str, text: str, title: str = \"\") -> list:\n",
+        "    \"\"\"Generate text embedding with a Large Language Model.\"\"\"\n",
+        "    model = TextEmbeddingModel.from_pretrained(model_name)\n",
+        "\n",
+        "    text_embedding_input = TextEmbeddingInput(\n",
+        "        task_type=task_type, title=title, text=text)\n",
+        "    embeddings = model.get_embeddings([text_embedding_input])\n",
+        "    return embeddings[0].values\n",
+        "\n",
+        "embedding = text_embedding(\n",
+        "    model_name=MODEL_NAME, task_type=TASK_TYPE, text=TEXT, title=TITLE)\n",
+        "print(len(embedding))"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "text_embedding_api_cloud_next_new_models.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -72,7 +72,9 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This notebook demonstrates finetuning a [JAX ViT-B16 model](https://github.com/google-research/vision_transformer#available-vit-models) for image classification task on GPU and deploying them on Vertex AI for online prediction."
+        "This notebook demonstrates finetuning a [JAX ViT-B16 model](https://github.com/google-research/vision_transformer#available-vit-models) for image classification task on GPU and deploying them on Vertex AI for online prediction.\n",
+        "\n",
+        "Learn more about [Generative AI Support in Vertex AI](https://cloud.google.com/blog/products/ai-machine-learning/vertex-ai-model-garden-and-generative-ai-studio)."
       ]
     },
     {
@@ -83,18 +85,21 @@
       "source": [
         "### Objective\n",
         "\n",
-        "In this tutorial, you learn how to:\n",
+        "In this tutorial, you learn how fine-tune, deploy and predict with a Vertex AI pretrained JAX Vision Transformer based model.\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "- Vertex AI Model Garden\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI Online Prediction\n",
+        "\n",
+        "The steps performed are:\n",
         "\n",
         "- Finetune a JAX Vision Transformer based model.\n",
         "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
         "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
-        "- Run online predictions for image classification.\n",
-        "\n",
-        "This tutorial uses the following Google Cloud ML services and resources:\n",
-        "\n",
-        "- Vertex AI Training\n",
-        "- Vertex AI Model Registry\n",
-        "- Vertex AI Online Prediction"
+        "- Run online predictions for image classification.\n"
       ]
     },
     {
@@ -165,10 +170,10 @@
       "outputs": [],
       "source": [
         "# Automatically restart kernel after installs so that your environment can access the new packages.\n",
-        "import IPython\n",
+        "# import IPython\n",
         "\n",
-        "app = IPython.Application.instance()\n",
-        "app.kernel.do_shutdown(True)"
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
       ]
     },
     {
@@ -214,7 +219,7 @@
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
         "\n",
         "# Set the project id\n",
         "! gcloud config set project {PROJECT_ID}"
@@ -511,7 +516,7 @@
         "### Prepare dataset\n",
         "\n",
         "If you are not using [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview#all_datasets), then you need to prepare your dataset and store it on Cloud Storage. The following example shows\n",
-        "how to do this for the [tf_flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers). If using TensorFlow Datasets, you can just pass\n",
+        "how to do this for the [tf_flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers). If using TensorFlow Datasets, you pass\n",
         "the dataset name such as `tf_flowers` to the `--config.dataset` flag and bypass this section."
       ]
     },
@@ -853,7 +858,7 @@
         "jax_vit_model.delete()\n",
         "\n",
         "# Delete Cloud Storage objects that were created.\n",
-        "delete_bucket = False\n",
+        "delete_bucket = True\n",
         "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
         "    ! gsutil -m rm -r $BUCKET_URI"
       ]

notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -282,18 +282,21 @@
         "from io import BytesIO\n",
         "\n",
         "import matplotlib.pyplot as plt\n",
+        "from google.cloud import storage\n",
         "from PIL import Image\n",
         "\n",
+        "GCS_URI_PREFIX = \"gs://\"\n",
+        "\n",
         "# Training constants.\n",
         "TRAINING_JOB_PREFIX = \"train\"\n",
-        "TRAIN_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
+        "TRAIN_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
         "TRAIN_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
         "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
         "TRAIN_NUM_GPU = 1\n",
         "RESOLUTION = 512\n",
         "\n",
         "# Prediction constants.\n",
-        "PREDICTION_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
+        "PREDICTION_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
         "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
         "DEPLOY_JOB_PREFIX = \"deploy\"\n",
@@ -317,6 +320,21 @@
         "    return gcs_path\n",
         "\n",
         "\n",
+        "def download_gcs_file_to_local(gcs_uri: str, local_path: str):\n",
+        "    \"\"\"Download a gcs file to a local path.\n",
+        "\n",
+        "    Args:\n",
+        "      gcs_uri: A string of file path on GCS.\n",
+        "      local_path: A string of local file path.\n",
+        "    \"\"\"\n",
+        "    if not gcs_uri.startswith(GCS_URI_PREFIX):\n",
+        "        raise ValueError(f\"{gcs_uri} is not a GCS path starting with {GCS_URI_PREFIX}.\")\n",
+        "    client = storage.Client()\n",
+        "    os.makedirs(os.path.dirname(local_path), exist_ok=True)\n",
+        "    with open(local_path, \"wb\") as f:\n",
+        "        client.download_blob_to_file(gcs_uri, f)\n",
+        "\n",
+        "\n",
         "def deploy_model(model_path, service_account):\n",
         "\n",
         "    deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX)\n",
@@ -420,7 +438,11 @@
         "from keras_cv.models import StableDiffusion\n",
         "\n",
         "model = StableDiffusion(img_height=RESOLUTION, img_width=RESOLUTION, jit_compile=True)\n",
-        "if model_path:\n",
+        "if model_path.startswith(GCS_URI_PREFIX):\n",
+        "    local_model_path = \"/tmp/saved_model.h5\"\n",
+        "    download_gcs_file_to_local(model_path, local_model_path)\n",
+        "    model.diffusion_model.load_weights(local_model_path)\n",
+        "elif model_path:\n",
         "    model.diffusion_model.load_weights(model_path)"
       ]
     },
@@ -568,7 +590,7 @@
       },
       "source": [
         "## Finetune models\n",
-        "This section shows how to finetune Keras Stable diffusion models with trainig dockers.\n",
+        "This section shows how to finetune Keras Stable diffusion models with training dockers.\n",
         "\n",
         "If you would like to use finetuned models, please go to the section `Run inferences`."
       ]

notebooks/community/model_garden/model_garden_keras_yolov8.ipynb

@@ -0,0 +1,772 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden Keras YOLOv8\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_keras_yolov8.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_keras_yolov8.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_keras_yolov8.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9\n",
+        "\n",
+        "You can open this notebook directly in Colab, or create [google managed](https://cloud.google.com/vertex-ai/docs/workbench/managed/create-instance) or [user managed](https://cloud.google.com/vertex-ai/docs/workbench/user-managed/create-new) Workbench instances."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [Keras YOLOv8](https://keras.io/api/keras_cv/models/tasks/yolo_v8_detector/) in Vertex AI Model Garden."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0z9r_mBmDeYh"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "* Run local inferences for pretrained or customized models\n",
+        "\n",
+        "* Deploy pretrained or customized models in Google Cloud Vertex AI\n",
+        "\n",
+        "* Finetune models in Google Cloud Vertex AI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "AEnkHABrDijz"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af989c0e437d"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "\n",
+        "The dataset used for this tutorial is the Salads category of the [OpenImages dataset](https://www.tensorflow.org/datasets/catalog/open_images_v4) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). This dataset does not require any feature engineering. The version of the dataset you will use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the bounding box locations and corresponding type of salad items in an image from a class of five items: Salad, Seafood, Tomato, Baked Goods, or Cheese."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    # Configs for Colab notebooks.\n",
+        "    ! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "\n",
+        "# Configs for all notebooks.\n",
+        "! pip3 install --quiet keras-cv==0.6.1\n",
+        "! pip3 install --quiet keras-core==0.1.0"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "### Set your project, region and buckets\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)\n",
+        "\n",
+        "You can change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n",
+        "\n",
+        "You can create a storage bucket to store intermediate artifacts such as datasets, trained models etc."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YjNCFxq0JxlA"
+      },
+      "outputs": [],
+      "source": [
+        "# The project and bucket are for experiments below.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# The form for BUCKET_URI is gs://<bucket-name>.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(STAGING_BUCKET, \"keras_yolov8\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "uDjp76aaLZY9"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5uv7-iDKLbO0"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZZFPe_GezXg8"
+      },
+      "source": [
+        "### Define constants and common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "import io\n",
+        "from datetime import datetime\n",
+        "from typing import Dict, List, Union\n",
+        "\n",
+        "import keras_cv\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "import yaml\n",
+        "from google.protobuf import json_format\n",
+        "from google.protobuf.struct_pb2 import Value\n",
+        "from keras_cv import visualization\n",
+        "from PIL import Image\n",
+        "\n",
+        "TRAIN_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAIN_NUM_GPU = 2\n",
+        "TRAIN_CONTAINER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-yolov8-train\"\n",
+        ")\n",
+        "TRAINING_JOB_PREFIX = \"train_yolov8\"\n",
+        "\n",
+        "UPLOAD_JOB_PREFIX = \"upload_yolov8\"\n",
+        "DEPLOY_JOB_PREFIX = \"deploy_yolov8\"\n",
+        "SERVING_CONTAINER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest\"\n",
+        ")\n",
+        "SERVING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "SERVING_MACHINE_TYPE = \"n1-standard-4\"\n",
+        "SERVING_CONTAINER_ARGS = [\"--allow_precompilation\", \"--allow_compression\"]\n",
+        "\n",
+        "RESOLUTION = 512\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str):\n",
+        "    \"\"\"Generates a job name with date time when triggering training or deployment\n",
+        "    jobs in Vertex AI.\n",
+        "    \"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def load_img(path):\n",
+        "    \"\"\"Reads image from path and return PIL.Image instance.\"\"\"\n",
+        "    img = tf.io.read_file(path)\n",
+        "    img = tf.image.decode_jpeg(img, channels=3)\n",
+        "    return Image.fromarray(np.uint8(img)).convert(\"RGB\")\n",
+        "\n",
+        "\n",
+        "def decode_image(image_str_tensor: tf.string) -> tf.float32:\n",
+        "    \"\"\"Converts and resizes image bytes to image tensor.\"\"\"\n",
+        "    image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)\n",
+        "    image = tf.image.resize(image, (RESOLUTION, RESOLUTION))\n",
+        "    return image\n",
+        "\n",
+        "\n",
+        "def get_label_map(label_map_yaml_filepath):\n",
+        "    \"\"\"Returns class id to label mapping given a filepath to the label map.\"\"\"\n",
+        "    with tf.io.gfile.GFile(label_map_yaml_filepath, \"rb\") as input_file:\n",
+        "        label_map = yaml.safe_load(input_file.read())[\"label_map\"]\n",
+        "    return label_map\n",
+        "\n",
+        "\n",
+        "def get_prediction_instances(test_filepath, new_width=-1):\n",
+        "    \"\"\"Generate instance from image path to pass to Vertex AI Endpoint for prediction.\"\"\"\n",
+        "    if new_width <= 0:\n",
+        "        with tf.io.gfile.GFile(test_filepath, \"rb\") as input_file:\n",
+        "            encoded_string = base64.b64encode(input_file.read()).decode(\"utf-8\")\n",
+        "    else:\n",
+        "        img = load_img(test_filepath)\n",
+        "        width, height = img.size\n",
+        "        print(\"original input image size: \", width, \" , \", height)\n",
+        "        new_height = int(height * new_width / width)\n",
+        "        new_img = img.resize((new_width, new_height))\n",
+        "        print(\"resized input image size: \", new_width, \" , \", new_height)\n",
+        "        buffered = io.BytesIO()\n",
+        "        new_img.save(buffered, format=\"JPEG\")\n",
+        "        encoded_string = base64.b64encode(buffered.getvalue()).decode(\"utf-8\")\n",
+        "\n",
+        "    instances = [\n",
+        "        {\n",
+        "            \"encoded_image\": {\"b64\": encoded_string},\n",
+        "        }\n",
+        "    ]\n",
+        "    return instances\n",
+        "\n",
+        "\n",
+        "def predict_custom_trained_model(\n",
+        "    project: str,\n",
+        "    endpoint_id: str,\n",
+        "    instances: Union[Dict, List[Dict]],\n",
+        "    location: str = \"us-central1\",\n",
+        "):\n",
+        "    # The AI Platform services require regional API endpoints.\n",
+        "    client_options = {\"api_endpoint\": f\"{location}-aiplatform.googleapis.com\"}\n",
+        "    # Initialize client that will be used to create and send requests.\n",
+        "    # This client only needs to be created once, and can be reused for multiple requests.\n",
+        "    client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)\n",
+        "    parameters_dict = {}\n",
+        "    parameters = json_format.ParseDict(parameters_dict, Value())\n",
+        "    endpoint = client.endpoint_path(\n",
+        "        project=project, location=location, endpoint=endpoint_id\n",
+        "    )\n",
+        "    response = client.predict(\n",
+        "        endpoint=endpoint, instances=instances, parameters=parameters\n",
+        "    )\n",
+        "    return response.predictions, response.deployed_model_id"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "epo-RHXzcBBT"
+      },
+      "source": [
+        "## Run local inferences with pretrained model\n",
+        "\n",
+        "This section shows how to run inferences locally with YOLOv8-M pretrained on PascalVOC 2012 object detection task, which consists of 20 classes.\n",
+        "\n",
+        "Load image from Cloud Storage and decode as Tensor."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6zsa9vnBHhvO"
+      },
+      "outputs": [],
+      "source": [
+        "test_filepath = \"\"  # @param {type:\"string\"}\n",
+        "img_bytes = tf.io.read_file(test_filepath)\n",
+        "image = tf.expand_dims(decode_image(img_bytes), axis=0)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2wC-pSYR0jjU"
+      },
+      "source": [
+        "Load model pretrained on PascalVOC 2012."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7nvPEly_4Vm6"
+      },
+      "outputs": [],
+      "source": [
+        "model = keras_cv.models.YOLOV8Detector.from_preset(\n",
+        "    \"yolo_v8_m_pascalvoc\",\n",
+        "    bounding_box_format=\"xywh\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZrijGrxT0lvC"
+      },
+      "source": [
+        "Then run inferences and visualize results."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "65yEa4N0xcTS"
+      },
+      "outputs": [],
+      "source": [
+        "decoded = model.predict(image)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "n8-X3gA5xV_l"
+      },
+      "outputs": [],
+      "source": [
+        "# Classes in PascalVOC 2012 dataset.\n",
+        "class_ids = [\n",
+        "    \"Aeroplane\",\n",
+        "    \"Bicycle\",\n",
+        "    \"Bird\",\n",
+        "    \"Boat\",\n",
+        "    \"Bottle\",\n",
+        "    \"Bus\",\n",
+        "    \"Car\",\n",
+        "    \"Cat\",\n",
+        "    \"Chair\",\n",
+        "    \"Cow\",\n",
+        "    \"Dining Table\",\n",
+        "    \"Dog\",\n",
+        "    \"Horse\",\n",
+        "    \"Motorbike\",\n",
+        "    \"Person\",\n",
+        "    \"Potted Plant\",\n",
+        "    \"Sheep\",\n",
+        "    \"Sofa\",\n",
+        "    \"Train\",\n",
+        "    \"Tvmonitor\",\n",
+        "    \"Total\",\n",
+        "]\n",
+        "class_mapping = dict(zip(range(len(class_ids)), class_ids))\n",
+        "\n",
+        "visualization.plot_bounding_box_gallery(\n",
+        "    image,\n",
+        "    value_range=(0, 255),\n",
+        "    rows=1,\n",
+        "    cols=1,\n",
+        "    y_pred=decoded,\n",
+        "    scale=5,\n",
+        "    font_scale=0.7,\n",
+        "    bounding_box_format=\"xywh\",\n",
+        "    class_mapping=class_mapping,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RB_xY9ipr7ZU"
+      },
+      "source": [
+        "## Finetune models\n",
+        "This section shows how to finetune the Keras YOLOv8 model with training dockers and then deploy to Vertex AI Endpoint resource. The accepted dataset format is a CSV formatted as it would for [AutoML Image Object Detection](https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#input-files), without an `ML_USE` column."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pkNc7jyq1js1"
+      },
+      "outputs": [],
+      "source": [
+        "input_csv_path = \"gs://cloud-samples-data/vision/salads.csv\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ee7Hzq8O5jgF"
+      },
+      "source": [
+        "### Start training jobs\n",
+        "The following code block shows some of the possible hyperparameters that can be set. The settings are for demonstration purposes only. Parameters such as `batch_size`, `learning_rate`, and `epochs` be overridden when used. `backbone` must be one of the following:\n",
+        "* `yolo_v8_xs_backbone`\n",
+        "* `yolo_v8_s_backbone`\n",
+        "* `yolo_v8_m_backbone`\n",
+        "* `yolo_v8_l_backbone`\n",
+        "* `yolo_v8_xl_backbone`\n",
+        "* `yolo_v8_xs_backbone_coco`\n",
+        "* `yolo_v8_s_backbone_coco`\n",
+        "* `yolo_v8_m_backbone_coco`\n",
+        "* `yolo_v8_l_backbone_coco`\n",
+        "* `yolo_v8_xl_backbone_coco`\n",
+        "\n",
+        "If looking for a preset with pretrained weights, choose one of `yolo_v8_xs_backbone_coco`, `yolo_v8_s_backbone_coco`, `yolo_v8_m_backbone_coco`, `yolo_v8_l_backbone_coco`, `yolo_v8_xl_backbone_coco`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "# Hyperparameters\n",
+        "epochs = 10\n",
+        "learning_rate = 0.0005\n",
+        "fpn_depth = 3\n",
+        "confidence_threshold = 0.02\n",
+        "iou_threshold = 0.3\n",
+        "backbone = \"yolo_v8_xl_backbone_coco\"\n",
+        "\n",
+        "train_job_name = get_job_name_with_datetime(TRAINING_JOB_PREFIX)\n",
+        "model_dir = os.path.join(MODEL_BUCKET, train_job_name)\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"disk_spec\": {\n",
+        "            \"boot_disk_type\": \"pd-ssd\",\n",
+        "            \"boot_disk_size_gb\": 500,\n",
+        "        },\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
+        "            \"command\": [],\n",
+        "            \"env\": [\n",
+        "                {\n",
+        "                    \"name\": \"RESOLUTION\",\n",
+        "                    \"value\": f\"{RESOLUTION}\",\n",
+        "                },\n",
+        "            ],\n",
+        "            \"args\": [\n",
+        "                f\"--input_csv_path={input_csv_path}\",\n",
+        "                f\"--output_model_dir={model_dir}\",\n",
+        "                f\"--epochs={epochs}\",\n",
+        "                f\"--pretrained_backbone={backbone}\",\n",
+        "                f\"--fpn_depth={fpn_depth}\",\n",
+        "                f\"--learning_rate={learning_rate}\",\n",
+        "                f\"--confidence_threshold={confidence_threshold}\",\n",
+        "                f\"--iou_threshold={iou_threshold}\",\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "train_job = aiplatform.CustomJob(\n",
+        "    display_name=train_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "train_job.run()\n",
+        "\n",
+        "print(\"The trained model is saved in: \", model_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9KBJ0ySVYX47"
+      },
+      "source": [
+        "### Prediction\n",
+        "This section shows how to deploy and make online predictions with the model.\n",
+        "\n",
+        "1. Upload and deploy models\n",
+        "2. Run predictions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "K6rUSSKmYZJ6"
+      },
+      "outputs": [],
+      "source": [
+        "upload_job_name = get_job_name_with_datetime(UPLOAD_JOB_PREFIX)\n",
+        "\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=upload_job_name,\n",
+        "    artifact_uri=model_dir,\n",
+        "    serving_container_image_uri=SERVING_CONTAINER_URI,\n",
+        "    serving_container_args=SERVING_CONTAINER_ARGS,\n",
+        ")\n",
+        "\n",
+        "print(\"The uploaded model name is: \", upload_job_name)\n",
+        "\n",
+        "deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX)\n",
+        "\n",
+        "endpoint = model.deploy(\n",
+        "    deployed_model_display_name=deploy_model_name,\n",
+        "    machine_type=SERVING_MACHINE_TYPE,\n",
+        "    traffic_split={\"0\": 100},\n",
+        "    accelerator_type=SERVING_ACCELERATOR_TYPE,\n",
+        "    accelerator_count=1,\n",
+        "    min_replica_count=1,\n",
+        "    max_replica_count=1,\n",
+        ")\n",
+        "print(\"The deployed job name is: \", deploy_model_name)\n",
+        "\n",
+        "endpoint_id = endpoint.name\n",
+        "print(\"endpoint id is: \", endpoint_id)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a879effaf402"
+      },
+      "source": [
+        "Load image from Cloud Storage, resize, and encode."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VDznWEMmbwj4"
+      },
+      "outputs": [],
+      "source": [
+        "test_filepath = \"gs://cloud-ml-data/img/openimage/1302/4677521502_6f2767039c_o.jpg\"  # @param {type:\"string\"}\n",
+        "image_bytes = tf.io.read_file(test_filepath)\n",
+        "image_resized = tf.expand_dims(decode_image(image_bytes), axis=0)\n",
+        "\n",
+        "instances = get_prediction_instances(test_filepath, new_width=640)\n",
+        "\n",
+        "predictions, _ = predict_custom_trained_model(\n",
+        "    project=PROJECT_ID, location=REGION, endpoint_id=endpoint_id, instances=instances\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "14e889492871"
+      },
+      "source": [
+        "Run online predictions using the endpoint and visualize the result."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2bx1cW0IdXqp"
+      },
+      "outputs": [],
+      "source": [
+        "predictions_dict = {\n",
+        "    \"boxes\": tf.expand_dims(predictions[0][\"boxes\"], axis=0),\n",
+        "    \"classes\": tf.expand_dims(predictions[0][\"classes\"], axis=0),\n",
+        "    \"confidence\": tf.expand_dims(predictions[0][\"confidence\"], axis=0),\n",
+        "    \"num_detections\": predictions[0][\"num_detections\"],\n",
+        "}\n",
+        "\n",
+        "label_map = get_label_map(os.path.join(model_dir, \"label_map.yaml\"))\n",
+        "\n",
+        "visualization.plot_bounding_box_gallery(\n",
+        "    image_resized,\n",
+        "    value_range=(0, 255),\n",
+        "    rows=1,\n",
+        "    cols=1,\n",
+        "    y_pred=predictions_dict,\n",
+        "    scale=5,\n",
+        "    font_scale=0.7,\n",
+        "    bounding_box_format=\"xywh\",\n",
+        "    class_mapping=label_map,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "### Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Deletes custom train jobs.\n",
+        "train_job.delete()\n",
+        "# Undeploys models and deletes endpoints.\n",
+        "endpoint.delete(force=True)\n",
+        "# Deletes models.\n",
+        "model.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1dijQDiZWegt"
+      },
+      "source": [
+        "## References\n",
+        "\n",
+        "- [Efficient Object Detection with YOLOV8 and KerasCV](https://keras.io/examples/vision/yolov8/)\n",
+        "- [Keras YOLOv8 API Documentation](https://keras.io/api/keras_cv/models/tasks/yolo_v8_detector/)\n",
+        "- [Keras YOLOv8 Backbones](https://keras.io/api/keras_cv/models/backbones/yolo_v8/)"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_keras_yolov8.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb

@@ -0,0 +1,580 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MediaPipe with Face Stylizer\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9\n",
+        "\n",
+        "**_NOTE_**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to customize an on-device face stylizer model in Vertex AI Model Garden.\n",
+        "\n",
+        "The MediaPipe face stylizer solution provides several models you can use immediately to transform the face to the styles including (cartoon, oil painting, etc.) in your application. However, if you need to transfer the face to an unseen style not covered by the provided models, you can customize the pretrained model with your own data and MediaPipe Model Maker. This model modification tool fine-tune a portion of the model using data you provide. This method is faster than training a new model from scatch and can produce a model adapt to your specific application.\n",
+        "\n",
+        "The following sections show you how to use Model Maker to retrain a pre-built model for face stylization with your own data on Vertex AI, which you can then use with the MediaPipe Face Stylizer.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Customize a Face Stylizer model\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to customize new models\n",
+        "  * Export customized models\n",
+        "\n",
+        "* Cleanup resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade pip\n",
+        "\n",
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
+        "\n",
+        "\n",
+        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
+        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
+        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"model.tflite\"\n",
+        ")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define training machine specs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_face_stylizer_%s\" % now\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAINING_ACCELERATOR_COUNT = 2"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-rsdAcBV-vlf"
+      },
+      "source": [
+        "## Train your customized models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "LmZ4efOd-sak"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Retraining the face stylizer model requires user to provide a single stylized face image. The stylized face is expected to be forward facing with visible left right eyes and mouth. The face should only have minor rotation, i.e. less than 30 degress around the yaw, pitch, and roll axes.\n",
+        "\n",
+        "You can upload an image to Google Cloud Storage or use our [provided example](https://storage.googleapis.com/mediapipe-assets/face_stylizer_style_color_sketch.jpg)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "training_data_path = \"gs://mediapipe-assets/face_stylizer_style_color_sketch.jpg\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "# Retrain model\n",
+        "\n",
+        "Once you have provided an input image, you can begin retraining the face stylizer model to adapt to the new style. This type of model modification is called transfer learning. The instructions below use the data prepared in the previous section to retrain a face stylizer model to apply cartoon style to the raw human face.\n",
+        "\n",
+        "**_NOTE_**: For this type of model, the retraining process causes the model to forget any style it can apply before. Once the retraining is complete, the new model can only apply the new style defined by the new stylized image.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kySYKb3-hnr6"
+      },
+      "source": [
+        "## Set retraining options\n",
+        "There are a few required settings to run a retraining aside from your training dataset:\n",
+        "\n",
+        "* **Swap layers:** The `swap_layers` parameter is used to determine how to mix the latent code layers between the learned style and the raw face images. The latent code is represented as a tensor of shape [1, 12, 512]. The second dimension of the latent code tensor is called the layer. The face stylizer mixes the learned style and raw face images by generating a weighted sum of the two latent codes on the swap layers. The swap layers are therefore integers within [1, 12]. The more layers are set, the more style will be applied to the output image. Although there is no explicit mapping between the style semantics and the layer index, the shallow layers, e.g. 8, 9, represent the global features of the face, while the deep layers, e.g. 10, 11, represent the fine-grained features of the face. The output stylized image is sensitive to the setting of swap layers. By default, it is set to [8, 9, 10, 11].\n",
+        "* **Learning rate and epochs:** Use `learning_rate` and epochs` to specify the these two hyperparameters. learning_rate is set to 4e-4 by default. epochs defines the number of iterations to fine-tune the BlazeStyleGAN model and are set to 100 by default. The lower the learning rate is, the greater the epochs is expected to retrain the model to converge.\n",
+        "* **Batch size:** The `batch_size` is used to define the number of latent code samples we sample around the latent code extracted by the encoder with the input image. The batch of latent codes are used to fine-tune the decoder. The greater the batch size usually yield to better performance. It is also limited by the hardware memory. For A100 GPU, the maximum batch size is 8. For P100 and T4 GPU, the maximum batch size is 2.\n",
+        "\n",
+        "Further more advanced parameters that you can configure are `alpha`, `perception_loss_weight`, `adv_loss_weight`, `beta_1` and `beta_2`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "# The layers of feature to be interpolated between encoding features and\n",
+        "# StyleGAN input features.\n",
+        "swap_layers: str = \"[8, 9, 10, 11]\"  # @param {type:\"string\"}\n",
+        "# The learning rate to use for gradient descent training.\n",
+        "learning_rate: float = 0.0001  # @param {type:\"number\"}\n",
+        "# Number of training iterations over the dataset.\n",
+        "epochs: int = 100  # @param {type:\"slider\", min:0, max:100, step:1}\n",
+        "# Batch size for training.\n",
+        "batch_size: int = 2  # @param {type:\"number\"}\n",
+        "\n",
+        "\n",
+        "# Other supported options\n",
+        "\n",
+        "# Weighting coefficient of style latent for swapping layer interpolation.\n",
+        "# Its valid range is [0, 1]. The greater weight means stronger style is\n",
+        "# applied to the output image. Expect to set it to a small value,\n",
+        "# i.e. < 0.1.\n",
+        "alpha: float = 0.1  # @param {type:\"number\"}\n",
+        "\n",
+        "# Weighting coefficients of image perception quality loss. It contains three\n",
+        "# coefficients, l1, content, and style which control the difference between the\n",
+        "# generated image and raw input image, the content difference between generated\n",
+        "# face and raw input face, and the how similar the style between the generated\n",
+        "# image and raw input image. Users can increase the style weight to enforce\n",
+        "# stronger style or the content weight to reserve more raw input face details.\n",
+        "# Weight for L1 loss.\n",
+        "perception_loss_l1: float = 0.5  # @param {type:\"number\"}\n",
+        "# Weight for content loss.\n",
+        "perception_loss_content: float = 4.0  # @param {type:\"number\"}\n",
+        "# Weight for stlye loss.\n",
+        "perception_loss_style: float = 1.0  # @param {type:\"number\"}\n",
+        "\n",
+        "# Weighting coeffcieint of adversarial loss versus image perceptual quality loss.\n",
+        "# This hyperparameter is used to control the realism of the generated image. It\n",
+        "# expects a small value, i.e. < 0.2.\n",
+        "adv_loss_weight: float = 0.2  # @param {type:\"number\"}\n",
+        "# beta_1 used in tf.keras.optimizers.Adam.\n",
+        "beta_1: float = 0.0  # @param {type:\"number\"}\n",
+        "# beta_2 used in tf.keras.optimizers.Adam.\n",
+        "beta_2: float = 0.99  # @param {type:\"number\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "### Run retraining\n",
+        "With your training dataset and retraining options prepared, you are ready to start the retraining process. This process requires running on GPU and can take a few minutes to a few hours depending on your available compute resources. On Vertex AI with GPU processing, the example retraining below takes about 2 minutes.\n",
+        "\n",
+        "To begin the fine-tuning process, use the following code:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAINING_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--task_name=face_stylizer\",\n",
+        "                \"--training_data_path=%s\" % training_data_path,\n",
+        "                \"--model_export_path=%s\" % model_export_path,\n",
+        "                \"--evaluation_result_path=%s\" % model_export_path,\n",
+        "                \"--hparams=%s\"\n",
+        "                % json.dumps(\n",
+        "                    {\n",
+        "                        \"learning_rate\": learning_rate,\n",
+        "                        \"batch_size\": batch_size,\n",
+        "                        \"epochs\": epochs,\n",
+        "                        \"beta_1\": beta_1,\n",
+        "                        \"beta_2\": beta_2,\n",
+        "                    }\n",
+        "                ),\n",
+        "                \"--model_options=%s\"\n",
+        "                % json.dumps(\n",
+        "                    {\n",
+        "                        \"swap_layers\": json.loads(swap_layers),\n",
+        "                        \"alpha\": alpha,\n",
+        "                        \"perception_loss_l1\": perception_loss_l1,\n",
+        "                        \"perception_loss_content\": perception_loss_content,\n",
+        "                        \"perception_loss_style\": perception_loss_style,\n",
+        "                        \"adv_loss_weight\": adv_loss_weight,\n",
+        "                    }\n",
+        "                ),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "training_job = aiplatform.CustomJob(\n",
+        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "training_job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "## Export model\n",
+        "After retraining the model, you can save the Tensorflow Lite model and integrate it with your on-device application by following the [Face stylization task guide](https://developers.google.com/mediapipe/solutions/vision/face_stylizer)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "def copy_model(model_source, model_dest):\n",
+        "    ! gsutil cp {model_source} {model_dest}\n",
+        "\n",
+        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"face_stylizer.task\")\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import files\n",
+        "\n",
+        "    files.download(\"face_stylizer.task\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete training data and jobs.\n",
+        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
+        "    training_job.delete()\n",
+        "\n",
+        "!gsutil rm -r {STAGING_BUCKET}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mediapipe_face_stylizer.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb

@@ -0,0 +1,643 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MediaPipe with gesture recognition\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>                                                                                               <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9\n",
+        "\n",
+        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to train an on-device gesture recognition model in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train new models\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to train new models\n",
+        "  * Export models\n",
+        "\n",
+        "* Cleanup resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade pip\n",
+        "\n",
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "import tensorflow\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
+        "\n",
+        "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
+        "EVALUATION_RESULT_OUTPUT_FILE = os.path.join(\n",
+        "    EVALUATION_RESULT_OUTPUT_DIRECTORY, \"evaluation.json\"\n",
+        ")\n",
+        "\n",
+        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
+        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
+        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"gesture_recognizer.task\"\n",
+        ")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define training machine specs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_gesture_recognizer_%s\" % now\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAINING_ACCELERATOR_COUNT = 2"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-rsdAcBV-vlf"
+      },
+      "source": [
+        "## Train your customized models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Finetuning a model for gesture recognition requires a dataset with a directory structure following the pattern `<dataset_path>/<label_name>/<img_name>.*` (e.g. `my_custom_dataset/thumbs_up/img12.jpg`). In addition, one of the label names must be none. The none label represents any gesture that isn't classified as one of the other gestures.\n",
+        "\n",
+        "This example uses a rock paper scissors dataset sample which is available on Cloud Storage.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "training_data_path = (\n",
+        "    \"gs://mediapipe-tasks/gesture_recognizer/rps_data_sample\"  # @param {type:\"string\"}\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ci4IV6vdXRMD"
+      },
+      "source": [
+        "When Model Maker loads the dataset, it runs the pre-packaged hand detection model from MediaPipe Hands to detect the hand landmarks from the images. Any images without detected hands are ommitted from the dataset. The resulting dataset will contain the extracted hand landmark positions from each image, rather than images themselves.\n",
+        "\n",
+        "You can configure a few options that determine how the dataset is loaded:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aNHLSyFtXP7I"
+      },
+      "outputs": [],
+      "source": [
+        "# A boolean controlling whether to shuffle the dataset. Defaults to true.\n",
+        "shuffle = True  # @param {type:\"boolean\"}\n",
+        "# A float between 0 and 1 controlling the confidence threshold for hand detection\n",
+        "min_detection_confidence = 0.6  # @param {type:\"number\"}\n",
+        "# Configures how to split the dataset between training, validation and test data. Must sum to up 1.\n",
+        "split_ratio = \"0.8,0.1,0.1\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Set fine-tuning options\n",
+        "\n",
+        "You can customize the model using the by specifying ModelOptions and HParams. The ModelOptions contain parameters related to the model itself, while the HParams  contains parameters related to training and saving the model.\n",
+        "\n",
+        "The ModelOptions contain these customizable parameter that affects accuracy:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bDxsEaoGcibW"
+      },
+      "outputs": [],
+      "source": [
+        "# The fraction of the input units to drop. Used in dropout layer.\n",
+        "dropout_rate: float = 0.05  # @param {type:\"number\"}\n",
+        "# A list of hidden layer widths for the gesture model. Each element\n",
+        "# in the list will create a new hidden layer with the specified width.\n",
+        "# The hidden layers are separated with BatchNorm, Dropout, and ReLU.\n",
+        "layer_widths: str = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fk0TTZbDdJPX"
+      },
+      "source": [
+        "HParams has the following list of customizable parameters which affect model accuracy:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "# The learning rate to use for gradient descent training.\n",
+        "learning_rate: float = 0.001  # @param {type:\"number\"}\n",
+        "# Batch size for training.\n",
+        "batch_size: int = 2  # @param {type:\"number\"}\n",
+        "# Number of training iterations over the dataset.\n",
+        "epochs: int = 10  # @param {type:\"slider\", min:0, max:100, step:1}\n",
+        "# An optional integer that indicates the number of training steps per\n",
+        "# epoch. If set to 0, the training pipeline calculates the default\n",
+        "# steps per epoch as the training dataset size divided by batch size.\n",
+        "steps_per_epoch: int = 0  # @param {type:\"number\"}\n",
+        "# Whether to shuffle the dataset before training\n",
+        "shuffle: bool = False  # @param {type:\"boolean\"}\n",
+        "# Learning rate decay to use for gradient descent training.\n",
+        "lr_decay: float = 0.99  # @param {type:\"number\"}\n",
+        "# Gamma parameter for focal loss. Defaults to 2\n",
+        "gamma: float = 2  # @param {type:\"number\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "### Run fine-tuning\n",
+        "With your training dataset and fine-tuning options prepared, you are ready to start the fine-tuning process. This process is resource intensive and can take a few minutes to complete. On Vertex AI with GPU processing, the example fine-tuning below takes between 1-2 minutes to train on approximately 500 images.\n",
+        "\n",
+        "To begin the fine-tuning process, use the following code:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
+        "evaluation_result_path = EVALUATION_RESULT_OUTPUT_DIRECTORY\n",
+        "\n",
+        "model_options = {\"dropout_rate\": dropout_rate}\n",
+        "if layer_widths:\n",
+        "    model_options[\"layer_widths\"] = layer_widths\n",
+        "\n",
+        "hparams = {\n",
+        "    \"learning_rate\": learning_rate,\n",
+        "    \"batch_size\": batch_size,\n",
+        "    \"epochs\": epochs,\n",
+        "    \"shuffle\": shuffle,\n",
+        "    \"lr_decay\": lr_decay,\n",
+        "    \"gamma\": gamma,\n",
+        "}\n",
+        "if steps_per_epoch:\n",
+        "    hparams[\"steps_per_epoch\"] = steps_per_epoch\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAINING_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--task_name=gesture_recognizer\",\n",
+        "                \"--training_data_path=%s\" % training_data_path,\n",
+        "                \"--model_export_path=%s\" % model_export_path,\n",
+        "                \"--evaluation_result_path=%s\" % evaluation_result_path,\n",
+        "                \"--split_ratio=%s\" % split_ratio,\n",
+        "                \"--model_options=%s\" % json.dumps(model_options),\n",
+        "                \"--hparams=%s\" % json.dumps(hparams),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "training_job = aiplatform.CustomJob(\n",
+        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "training_job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rXMF2tnV_WS0"
+      },
+      "source": [
+        "## Evaluate and export model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mV-Djz-frBni"
+      },
+      "source": [
+        "### Evaluate performance\n",
+        "\n",
+        "After fine-tuning the model, we evaluate the training result on a test dataset, which is typically a portion of your original dataset not used during training. Accuracy levels between 0.8 and 0.9 are generally considered very good, but your use case requirements may differ. You should also consider how fast the model can produce an inference. Higher accuracy frequently comes at the cost of longer inference times.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "09Rz1AYspK19"
+      },
+      "outputs": [],
+      "source": [
+        "def get_evaluation_result(evaluation_result_path):\n",
+        "    try:\n",
+        "        with tensorflow.io.gfile.GFile(evaluation_result_path, \"r\") as input_file:\n",
+        "            evalutation_result = json.loads(input_file.read())\n",
+        "        return evalutation_result[\"accuracy\"], evalutation_result[\"loss\"]\n",
+        "    except:\n",
+        "        print(\n",
+        "            \"Evaluation result not found. Your test dataset is likely \"\n",
+        "            + \"empty. You can adjust the size of your test dataset or adjust \"\n",
+        "            + \"how you split your dataset.\"\n",
+        "        )\n",
+        "        return None\n",
+        "\n",
+        "\n",
+        "evaluation_result = get_evaluation_result(EVALUATION_RESULT_OUTPUT_FILE)\n",
+        "\n",
+        "if evaluation_result is not None:\n",
+        "    print(\"Accuracy:\", evaluation_result[0])\n",
+        "    print(\"Loss:\", evaluation_result[1])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "### Export model\n",
+        "After finetuning and evaluating the model, you can save the Tensorflow Lite model, try it out in the [Gesture Recognizer](https://mediapipe-studio.webapps.google.com/demo/gesture_recognizer) demo in MediaPipe Studio or integrate it with your on-device application by following the [Gesture recognizer task guide](https://developers.google.com/mediapipe/solutions/vision/gesture_recognizer). The exported model contains the generates required model metadata, as well as a classification label file."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "\n",
+        "def copy_model(model_source, model_dest):\n",
+        "    ! gsutil cp {model_source} {model_dest}\n",
+        "\n",
+        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"gesture_recognizer.task\")\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import files\n",
+        "\n",
+        "    files.download(\"gesture_recognizer.task\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete training data and jobs.\n",
+        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
+        "    training_job.delete()\n",
+        "\n",
+        "!gsutil rm -r {STAGING_BUCKET}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mediapipe_gesture_recognition.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb

@@ -29,7 +29,7 @@
         "id": "TirJ-SGQseby"
       },
       "source": [
-        "# Vertex AI Model Garden MediaPipe With Image Classification\n",
+        "# Vertex AI Model Garden MediaPipe with image classification\n",
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -61,7 +61,9 @@
       "source": [
         "**_NOTE_**: This notebook has been tested in the following environment:\n",
         "\n",
-        "* Python version = 3.9"
+        "* Python version = 3.9\n",
+        "\n",
+        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
       ]
     },
     {
@@ -145,24 +147,132 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "WReHDGG5g0XY"
       },
       "source": [
-        "### Set up your Google Cloud project\n",
+        "#### Set your project ID\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
         "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
         "\n",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
         "\n",
-        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
-        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
         "\n",
-        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "import tensorflow\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
         "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+        "Initialize the Vertex AI SDK for Python for your project."
       ]
     },
     {
@@ -173,32 +283,8 @@
       },
       "outputs": [],
       "source": [
-        "import json\n",
-        "import os\n",
-        "from datetime import datetime\n",
-        "\n",
-        "import tensorflow\n",
-        "from google.cloud import aiplatform\n",
-        "\n",
         "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
         "\n",
-        "# The project and bucket are for experiments below.\n",
-        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
-        "# The form for BUCKET_URI is gs://<bucket-name>.\\n\",\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# You can choose a region from https://cloud.google.com/about/locations.\n",
-        "# Only regions prefixed by \"us\", \"asia\", or \"europe\" are supported.\n",
-        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
-        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
-        "assert REGION_PREFIX in (\n",
-        "    \"us\",\n",
-        "    \"europe\",\n",
-        "    \"asia\",\n",
-        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'\n",
-        "\n",
-        "! gcloud config set project $PROJECT_ID\n",
-        "\n",
         "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
         "\n",
         "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
@@ -234,7 +320,7 @@
         "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_image_classifier_%s\" % now\n",
         "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
         "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
-        "TRAINING_ACCELARATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "TRAINING_ACCELERATOR_COUNT = 2"
       ]
     },
@@ -280,14 +366,14 @@
       "source": [
         "### Set fine-tuning options\n",
         "\n",
-        "There are a few required settings to run fine-tuning aside from your training dataset: output directory for the model and the model architecture. Use HParams object export_dir parameter to specify a model output directory. Use the SupportedModels class to specify the model architecture. The image classifier solution supports the following model architectures:\n",
+        "You can pick between different model architectures to further customize your training:\n",
         "\n",
         "*   MobileNet-V2\n",
         "*   EfficientNet-Lite0\n",
         "*   EfficientNet-Lite2\n",
         "*   EfficientNet-Lite4\n",
         "\n",
-        "To set the required parameters, use the following code:"
+        "To set the model architecture and other training parameters, adjust the following values:"
       ]
     },
     {
@@ -354,7 +440,7 @@
         "    {\n",
         "        \"machine_spec\": {\n",
         "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
-        "            \"accelerator_type\": TRAINING_ACCELARATOR_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
         "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
         "        },\n",
         "        \"replica_count\": 1,\n",

notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb

@@ -0,0 +1,821 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MediaPipe with image generation\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9\n",
+        "\n",
+        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to customize a text-to-image generator model in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train a custom text-to-image generator\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to train new a generator\n",
+        "  * Export models\n",
+        "  * Generate images\n",
+        "* Cleanup resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab Only\n",
+        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade pip\n",
+        "\n",
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6d87ae73686c"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
+        "\n",
+        "MODEL_EXPORT_PATH = os.path.join(STAGING_BUCKET, \"model\")\n",
+        "\n",
+        "IMAGE_EXPORT_PATH = os.path.join(STAGING_BUCKET, \"image\")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define training and serving constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_stable_diffusion_%s\" % now\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-stable-diffusion-train\"\n",
+        "TRAINING_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
+        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
+        "TRAINING_ACCELERATOR_COUNT = 1\n",
+        "\n",
+        "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        "PREDICTION_PORT = 7080\n",
+        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
+        "UPLOAD_MODEL_NAME = \"mediapipe_stable_diffusion_model_%s\" % now"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XDq9TiRUc7dV"
+      },
+      "source": [
+        "## Train your customize image generator"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Oj4XAhsn78as"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Customizing a model for image generation requires a dataset that contains sample pictures of the OBJECT INSTANCE that you want to use in the generation. These images can be downloaded from Google Cloud Storage. You must also provide a description that describes the object in these images.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "# Path to the training data folder.\n",
+        "training_data_path = \"gs://mediapipe-tasks/image_generator/teapot\"  # @param {type:\"string\"}\n",
+        "# An instance description of the training data.\n",
+        "training_data_prompt = \"A monadikos teapot\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Set training options\n",
+        "\n",
+        "The Image Generator comes with a set of pre-defined HParams settings that work best for specific situations. You should select a template that best matches your use case.\n",
+        "\n",
+        "If applicable, you can also use one of our pre-trained models for these templates. These can be used directly and without training:\n",
+        "* [Object - Berry Bowl](https://storage.googleapis.com/mediapipe-tasks/image_generator/object/pytorch_lora_weights.bin)\n",
+        "* [Face](https://storage.googleapis.com/mediapipe-tasks/image_generator/face/pytorch_lora_weights.bin)\n",
+        "* [Style](https://storage.googleapis.com/mediapipe-tasks/image_generator/style/pytorch_lora_weights.bin)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bj31hD4W02Ui"
+      },
+      "outputs": [],
+      "source": [
+        "template = \"\"  # @param [\"\", \"Face\", \"Object\", \"Style\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "lYF-4ydq0k0u"
+      },
+      "source": [
+        "To set custom training parameters, adjust the following values:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "# Parameters about training configuration\n",
+        "# The learning rate to use for gradient descent training.\n",
+        "learning_rate: float = 0.00001  # @param {type:\"number\"}\n",
+        "# Number of training steps. If set to 0, uses the default value.\n",
+        "num_train_steps: int = 0  # @param {type:\"integer\"}\n",
+        "# Save the checkpoint in every n steps.\n",
+        "save_checkpoints_every_n: int = 100  # @param {type:\"integer\"}\n",
+        "# Batch size for training.\n",
+        "batch_size: int = 1  # @param {type:\"integer\"}\n",
+        "\n",
+        "# Dataset-related parameters\n",
+        "# Whether to use random horizontal flip on data.\n",
+        "random_flip: bool = False  # @param {type:\"boolean\"}\n",
+        "# Whether to use random largest square crop.\n",
+        "random_crop: bool = False  # @param {type:\"boolean\"}\n",
+        "# Whether to distort the color of the image (jittering order is random).\n",
+        "random_color_jitter: bool = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Hyperparameters for LoRA tuning\n",
+        "# The rank in the low-rank matrices. If set to 0, uses the default value.\n",
+        "lora_rank: int = 0  # @param {type:\"integer\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BuMz_X97F6yr"
+      },
+      "source": [
+        "## Configure image generation\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5NwgyEvYGOTt"
+      },
+      "source": [
+        "### Generation prompt\n",
+        "\n",
+        "You need to specify the prompt to use for image generation"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2dSmBAauGEWw"
+      },
+      "outputs": [],
+      "source": [
+        "prompt: str = \"Two monadikos teapots on a table\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nvhSeCt1GbUy"
+      },
+      "source": [
+        "### Customize image generation\n",
+        "\n",
+        "You may customize the image generation."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "iwmsxMzDGhOO"
+      },
+      "outputs": [],
+      "source": [
+        "# Number of steps to run inference.\n",
+        "number_inference_steps: int = 50  # @param {type:\"integer\"}\n",
+        "# Classifier-free guidance weight to use during inference. Weight must be is >= 1.0.\n",
+        "guidance_scale: float = 7.5  # @param {type:\"number\"}\n",
+        "#  Number of generated images per prompt.\n",
+        "number_generated_images: int = 8  # @param {type:\"integer\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "## Tune the image generator with Lora\n",
+        "Tune the image generator with Lora and generate new images based on your prompt. This can take up to 10 minutes on Vertex AI with a A100 GPU.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "model_export_path = MODEL_EXPORT_PATH\n",
+        "image_export_path = IMAGE_EXPORT_PATH\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAINING_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--task_name=stable_diffusion\",\n",
+        "                \"--model_export_path=%s\" % model_export_path,\n",
+        "                \"--image_export_path=%s\" % image_export_path,\n",
+        "                \"--training_data_path=%s\" % training_data_path,\n",
+        "                \"--training_data_prompt='%s'\" % training_data_prompt,\n",
+        "                \"--prompt='%s'\" % prompt,\n",
+        "                \"--hparams=%s\"\n",
+        "                % json.dumps(\n",
+        "                    {\n",
+        "                        \"learning_rate\": learning_rate,\n",
+        "                        \"num_train_steps\": num_train_steps,\n",
+        "                        \"save_checkpoints_every_n\": save_checkpoints_every_n,\n",
+        "                        \"batch_size\": batch_size,\n",
+        "                        \"random_flip\": random_flip,\n",
+        "                        \"random_crop\": random_crop,\n",
+        "                        \"random_color_jitter\": random_color_jitter,\n",
+        "                        \"lora_rank\": lora_rank,\n",
+        "                    }\n",
+        "                ),\n",
+        "                \"--generator_hparams=%s\"\n",
+        "                % json.dumps(\n",
+        "                    {\n",
+        "                        \"number_inference_steps\": number_inference_steps,\n",
+        "                        \"guidance_scale\": guidance_scale,\n",
+        "                        \"number_generated_images\": number_generated_images,\n",
+        "                    }\n",
+        "                ),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "training_job = aiplatform.CustomJob(\n",
+        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "training_job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Jdm1dOgsMH9O"
+      },
+      "source": [
+        "## Download generated images\n",
+        "\n",
+        "You can download and preview the generated images."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xsrMHpRvMErS"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "import matplotlib.pyplot as plt\n",
+        "\n",
+        "\n",
+        "def copy_image(images_source, images_dest):\n",
+        "    os.makedirs(images_dest, exist_ok=True)\n",
+        "    ! gsutil cp -r {images_source}/* {images_dest}\n",
+        "\n",
+        "\n",
+        "local_image_path = \"./images/\"\n",
+        "copy_image(IMAGE_EXPORT_PATH, local_image_path)\n",
+        "\n",
+        "for filename in os.listdir(local_image_path):\n",
+        "    absolute_path = os.path.join(local_image_path, filename)\n",
+        "    if os.path.isfile(absolute_path):\n",
+        "        im = plt.imread(absolute_path)\n",
+        "        plt.imshow(im)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zcKzIa5QeIIU"
+      },
+      "source": [
+        "## Download model\n",
+        "\n",
+        "After fine-tuning and evaluating the model, you can download the model and checkpoints."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "\n",
+        "def copy_model(model_source, model_dest):\n",
+        "    os.makedirs(model_dest, exist_ok=True)\n",
+        "    ! gsutil -m cp -r {model_source}/* {model_dest}\n",
+        "\n",
+        "\n",
+        "local_model_path = \"/models\"\n",
+        "copy_model(MODEL_EXPORT_PATH, local_model_path)\n",
+        "\n",
+        "! tar czf models.tar.gz {local_model_path}/*\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import files\n",
+        "\n",
+        "    files.download(\"models.tar.gz\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-gqs4Kc4u6U3"
+      },
+      "source": [
+        "## Use trained models on Vertex AI\n",
+        "\n",
+        "This section shows the way to test with trained models.\n",
+        "1. Upload and deploy models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
+        "2. Get [online predictions](https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions) from the deployed model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Y_HV9vMbvF_6"
+      },
+      "source": [
+        "### Upload model to Vertex AI Model Registry"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "OL6wdXf-1rta"
+      },
+      "outputs": [],
+      "source": [
+        "serving_env = {\n",
+        "    \"TASK\": \"text-to-image-lora\",\n",
+        "    \"BASE_MODEL_ID\": \"runwayml/stable-diffusion-v1-5\",\n",
+        "    \"FINETUNED_LORA_MODEL_PATH\": MODEL_EXPORT_PATH,\n",
+        "}\n",
+        "\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=UPLOAD_MODEL_NAME,\n",
+        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
+        "    serving_container_ports=[PREDICTION_PORT],\n",
+        "    serving_container_predict_route=\"/predictions/peft_serving\",\n",
+        "    serving_container_health_route=\"/ping\",\n",
+        "    serving_container_environment_variables=serving_env,\n",
+        ")\n",
+        "\n",
+        "model.wait()\n",
+        "\n",
+        "print(\"The uploaded model name is: \", UPLOAD_MODEL_NAME)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NjXg3QHw2GCT"
+      },
+      "source": [
+        "### Deploy the uploaded model\n",
+        "\n",
+        "You will deploy models in Google Cloud Vertex AI. The default setting will use 1 V100 GPU for deployment.\n",
+        "\n",
+        "Please create a Service Account for serving with dockers if you do not have one yet.\n",
+        "\n",
+        "The model deployment will take around 1 minute to finish."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BLTA7nRcw57D"
+      },
+      "outputs": [],
+      "source": [
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "service_account = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "endpoint = aiplatform.Endpoint.create(display_name=f\"{UPLOAD_MODEL_NAME}-endpoint\")\n",
+        "model.deploy(\n",
+        "    endpoint=endpoint,\n",
+        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    accelerator_count=1,\n",
+        "    deploy_request_timeout=1800,\n",
+        "    service_account=service_account,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tMfwrH32_uIM"
+      },
+      "source": [
+        "The docker container still needs to download and load the model after the endpoint is created. Therefore, we recommend waiting for 3 extra minutes before proceeding to the next cell.\n",
+        "\n",
+        "Once deployed, you can send a batch of text prompts to the endpoint to generate images."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wr7wg9JQ_tjM"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "from io import BytesIO\n",
+        "\n",
+        "import matplotlib.pyplot as plt\n",
+        "from PIL import Image\n",
+        "\n",
+        "instances = [\n",
+        "    {\"prompt\": \"Two monadikos teapots on a table\"},\n",
+        "    {\"prompt\": \"Two monadikos teapots on the floor\"},\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances)\n",
+        "\n",
+        "plt.figure()\n",
+        "_, grid = plt.subplots(1, len(instances))\n",
+        "for cell, prediction in zip(grid, response.predictions):\n",
+        "    image = Image.open(BytesIO(base64.b64decode(prediction)))\n",
+        "    cell.imshow(image)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up\n",
+        "After the export is complete, you can delete your training job."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
+        "    training_job.delete()\n",
+        "# Undeploys models and deletes endpoints.\n",
+        "endpoint.delete(force=True)\n",
+        "model.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CkRKXILVJc7Q"
+      },
+      "source": [
+        "You can also remove the output data."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vdIZsAstJTcA"
+      },
+      "outputs": [],
+      "source": [
+        "!gsutil rm -r {STAGING_BUCKET}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mediapipe_image_generation.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb

@@ -29,7 +29,7 @@
         "id": "TirJ-SGQseby"
       },
       "source": [
-        "# Vertex AI Model Garden MediaPipe With Object Detection\n",
+        "# Vertex AI Model Garden MediaPipe with object detection\n",
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
@@ -61,7 +61,9 @@
       "source": [
         "**_NOTE_**: This notebook has been tested in the following environment:\n",
         "\n",
-        "* Python version = 3.9"
+        "* Python version = 3.9\n",
+        "\n",
+        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
       ]
     },
     {
@@ -112,7 +114,7 @@
         "id": "z__i0w0lCAsW"
       },
       "source": [
-        "### Colab Only\n",
+        "### Colab only\n",
         "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
       ]
     },
@@ -145,24 +147,132 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "WReHDGG5g0XY"
       },
       "source": [
-        "### Set up your Google Cloud project\n",
+        "#### Set your project ID\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
         "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
         "\n",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
         "\n",
-        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
-        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
         "\n",
-        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "import tensorflow\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
         "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+        "Initialize the Vertex AI SDK for Python for your project."
       ]
     },
     {
@@ -173,32 +283,8 @@
       },
       "outputs": [],
       "source": [
-        "import json\n",
-        "import os\n",
-        "from datetime import datetime\n",
-        "\n",
-        "import tensorflow\n",
-        "from google.cloud import aiplatform\n",
-        "\n",
         "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
         "\n",
-        "# The project and bucket are for experiments below.\n",
-        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
-        "# The form for BUCKET_URI is gs://<bucket-name>.\\n\",\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# You can choose a region from https://cloud.google.com/about/locations.\n",
-        "# Only regions prefixed by \"us\", \"asia\", or \"europe\" are supported.\n",
-        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
-        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
-        "assert REGION_PREFIX in (\n",
-        "    \"us\",\n",
-        "    \"europe\",\n",
-        "    \"asia\",\n",
-        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'\n",
-        "\n",
-        "! gcloud config set project $PROJECT_ID\n",
-        "\n",
         "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
         "\n",
         "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
@@ -234,7 +320,7 @@
         "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_object_detector_%s\" % now\n",
         "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
         "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
-        "TRAINING_ACCELARATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "TRAINING_ACCELERATOR_COUNT = 2"
       ]
     },
@@ -359,12 +445,12 @@
       "source": [
         "### Set fine-tuning options\n",
         "\n",
-        "There are a few required settings to run fine-tuning aside from your training dataset: output directory for the model, and the model architecture. Use HParams to specify the export_dir parameter for the output directory. Use the SupportedModels class to specify the model architecture. The object detector solution supports the following model architectures:\n",
+        "You can pick between different model architectures to further customize your training:\n",
         "\n",
         "*   MobileNet-V2\n",
         "*   MobileNet-MultiHW-AVG\n",
         "\n",
-        "To set the parameters, use the following code:"
+        "To set the model architecture and other training parameters, adjust the following values:"
       ]
     },
     {
@@ -433,7 +519,7 @@
         "    {\n",
         "        \"machine_spec\": {\n",
         "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
-        "            \"accelerator_type\": TRAINING_ACCELARATOR_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
         "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
         "        },\n",
         "        \"replica_count\": 1,\n",

notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb

@@ -0,0 +1,618 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MediaPipe with text classification\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9\n",
+        "\n",
+        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to train an on-device text classification model in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train new models\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to train new models\n",
+        "  * Export models\n",
+        "\n",
+        "* Cleanup resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade pip\n",
+        "\n",
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tTy1gX11kCJY"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
+        "\n",
+        "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
+        "EVALUATION_RESULT_OUTPUT_FILE = os.path.join(\n",
+        "    EVALUATION_RESULT_OUTPUT_DIRECTORY, \"evaluation.json\"\n",
+        ")\n",
+        "\n",
+        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
+        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
+        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"model.tflite\"\n",
+        ")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define training machine specs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_text_classifier_%s\" % now\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAINING_ACCELERATOR_COUNT = 2"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-rsdAcBV-vlf"
+      },
+      "source": [
+        "## Train your customized models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Get the Dataset\n",
+        "\n",
+        "The following code block uses the [SST-2](https://nlp.stanford.edu/sentiment/index.html) (Stanford Sentiment Treebank) dataset which contains 67,349 movie reviews for training and 872 movie reviews for testing. The dataset has two classes: positive and negative movie reviews. Positive reviews are labeled with 1 and negative reviews with 0.\n",
+        "\n",
+        "The SST-2 dataset is stored as a TSV file. The only difference between the TSV and CSV formats is that TSV uses a tab `\\t` character as its delimiter and CSV uses a comma `,`.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "training_data_path = (\n",
+        "    \"gs://mediapipe-tasks/text_classifier/SST-2/train.tsv\"  # @param {type:\"string\"}\n",
+        ")\n",
+        "validation_data_path = (\n",
+        "    \"gs://mediapipe-tasks/text_classifier/SST-2/dev.tsv\"  # @param {type:\"string\"}\n",
+        ")\n",
+        "\n",
+        "# The delimiter used in the dataset.\n",
+        "delimiter = \"\\t\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Character used to quote fields that contain special characters\n",
+        "# like the `delimiter`.\n",
+        "quotechar = \"\\t\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Sequence of keys for the CSV columns (represented as a comma\n",
+        "# separated list). If empty, the first row of the CSV file is used\n",
+        "# as the keys\n",
+        "fieldnames = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Column name for the input text.\n",
+        "text_column = \"sentence\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Column name for the labels.\n",
+        "label_column = \"label\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Set fine-tuning options\n",
+        "\n",
+        "You can pick between different model architectures to further customize your training:\n",
+        "\n",
+        "*   Average Word Embedding Model\n",
+        "*   BERT-classifier\n",
+        "\n",
+        "To set the model architecture and other training parameters, adjust the following values:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "model_architecture = (\n",
+        "    \"average_word_embedding\"  # @param [\"average_word_embedding\", \"mobilebert\"]\n",
+        ")\n",
+        "\n",
+        "# The learning rate to use for gradient descent-based\n",
+        "# optimizers. Defaults to 3e-5 for the BERT-based classifier\n",
+        "# and 0 for the average word-embedding classifier because\n",
+        "# it does not need such an optimizer.\n",
+        "learning_rate: float = 0.0  # @param {type:\"number\"}\n",
+        "\n",
+        "# Batch size for training. Defaults to 32 for the average\n",
+        "# word-embedding classifier and 48 for the BERT-based\n",
+        "# classifier.\n",
+        "batch_size: int = 48  # @param {type:\"number\"}\n",
+        "\n",
+        "# Number of training iterations over the dataset. Defaults\n",
+        "# to 10 for the average word-embedding classifier and 3\n",
+        "# for the BERT-based classifier.\n",
+        "epochs: int = 10  # @param {type:\"slider\", min:0, max:100, step:1}\n",
+        "\n",
+        "# An integer that indicates the number of training steps per\n",
+        "# epoch. If set to 0, the training pipeline calculates the\n",
+        "# default steps per epoch as the training dataset size\n",
+        "# divided by batch size.\n",
+        "steps_per_epoch: int = 0  # @param {type:\"number\"}\n",
+        "\n",
+        "# Controls whether the dataset is shuffled before training.\n",
+        "shuffle: bool = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Length of the sequence to feed into the model.\n",
+        "seq_len: int = 256  # @param {type:\"number\"}\n",
+        "\n",
+        "# Whether to convert all uppercase characters to lowercase\n",
+        "# during preprocessing.\n",
+        "do_lower_case: bool = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# The rate for dropout.\n",
+        "dropout_rate: float = 0.2  # @param {type:\"number\"}\n",
+        "\n",
+        "# Dimension of the word embedding. Only used for the Average Word\n",
+        "# Embedding Model.\n",
+        "wordvec_dim: int = 16  # @param {type:\"number\"}\n",
+        "\n",
+        "# Number of words to generate the vocabulary from data.\n",
+        "# Only used for the Average Word Embedding Model.\n",
+        "vocab_size: int = 10000  # @param {type:\"number\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "### Run fine-tuning\n",
+        "With your training dataset and fine-tuning options prepared, you are ready to start the fine-tuning process. This process is resource intensive and can take a few minutes to a few hours depending on the model archtiecture and your available compute resources. On Vertex AI with GPU processing, the example fine-tuning below takes between 2-3 minutes to train an Average Word Embedding Model on the SST-2 dataset.\n",
+        "\n",
+        "To begin the fine-tuning process, use the following code:\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
+        "evaluation_result_path = EVALUATION_RESULT_OUTPUT_DIRECTORY\n",
+        "\n",
+        "preprocessing_params = {\n",
+        "    \"text_column\": text_column,\n",
+        "    \"label_column\": label_column,\n",
+        "    \"delimiter\": delimiter,\n",
+        "    \"quotechar\": quotechar,\n",
+        "}\n",
+        "if fieldnames:\n",
+        "    preprocessing_params[\"fieldnames\"] = [\n",
+        "        fieldname.strip() for fieldname in fieldnames.split(\",\")\n",
+        "    ]\n",
+        "\n",
+        "hparams = {\n",
+        "    \"learning_rate\": learning_rate,\n",
+        "    \"batch_size\": batch_size,\n",
+        "    \"epochs\": epochs,\n",
+        "    \"shuffle\": shuffle,\n",
+        "}\n",
+        "if steps_per_epoch:\n",
+        "    hparams[\"steps_per_epoch\"] = steps_per_epoch\n",
+        "\n",
+        "model_options = {\n",
+        "    \"dropout_rate\": dropout_rate,\n",
+        "    \"wordvec_dim\": wordvec_dim,\n",
+        "    \"do_lower_case\": do_lower_case,\n",
+        "    \"vocab_size\": vocab_size,\n",
+        "    \"dropout_rate\": dropout_rate,\n",
+        "}\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAINING_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--task_name=text_classifier\",\n",
+        "                \"--training_data_path=%s\" % training_data_path,\n",
+        "                \"--validation_data_path=%s\" % validation_data_path,\n",
+        "                \"--evaluation_result_path=%s\" % evaluation_result_path,\n",
+        "                \"--model_export_path=%s\" % model_export_path,\n",
+        "                \"--model_architecture=%s\" % model_architecture,\n",
+        "                \"--preprocessing_params=%s\" % json.dumps(preprocessing_params),\n",
+        "                \"--hparams=%s\" % json.dumps(hparams),\n",
+        "                \"--model_options=%s\" % json.dumps(model_options),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "training_job = aiplatform.CustomJob(\n",
+        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "training_job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rXMF2tnV_WS0"
+      },
+      "source": [
+        "## Export model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "After finetuning, you can save the Tensorflow Lite model, try it out in the [Text Classification](https://mediapipe-studio.webapps.google.com/demo/text_classifier) demo in MediaPipe Studio or integrate it with your on-device application by following the [Text classification task guide](https://developers.google.com/mediapipe/solutions/text/text_classifier). The exported model contains the generates required model metadata, as well as a classification label file."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "\n",
+        "def copy_model(model_source, model_dest):\n",
+        "    ! gsutil cp {model_source} {model_dest}\n",
+        "\n",
+        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"text_classification_model.tflite\")\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import files\n",
+        "\n",
+        "    files.download(\"text_classification_model.tflite\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete training data and jobs.\n",
+        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
+        "    training_job.delete()\n",
+        "\n",
+        "!gsutil rm -r {STAGING_BUCKET}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mediapipe_text_classification.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb

@@ -0,0 +1,961 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MoViNet video action recognition\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>                                                                                               <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MoViNet](https://github.com/tensorflow/models/tree/master/official/projects/movinet) for video action recognition in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train new models\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [hyperparameter tuning jobs](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to train new models\n",
+        "  * Find and export best models\n",
+        "\n",
+        "* Test trained models\n",
+        "  * Upload models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
+        "  * Run batch predictions\n",
+        "\n",
+        "* Clean up resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab Only\n",
+        "Run the following commands for Colab or skip this section if you use Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for running batch predictions with the fine tuned model.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# The GCP project ID for experiments.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Bucket URI with gs:// prefix.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# You can choose a region from https://cloud.google.com/about/locations.\n",
+        "# Only regions prefixed by \"us\", \"asia\", or \"europe\" are supported.\n",
+        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "CHECKPOINT_BUCKET = os.path.join(BUCKET_URI, \"ckpt\")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Download config files.\n",
+        "CONFIG_DIR = os.path.join(BUCKET_URI, \"config\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "OBJECTIVE = \"var\"\n",
+        "\n",
+        "# Data converter constants.\n",
+        "DATA_CONVERTER_JOB_PREFIX = \"data_converter\"\n",
+        "DATA_CONVERTER_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/data-converter\"\n",
+        "DATA_CONVERTER_MACHINE_TYPE = \"n1-highmem-8\"\n",
+        "IMAGE_SIZES = {\n",
+        "    \"a0\": 172,\n",
+        "    \"a1\": 172,\n",
+        "    \"a2\": 224,\n",
+        "    \"a3\": 256,\n",
+        "    \"a4\": 290,\n",
+        "    \"a5\": 320,\n",
+        "}\n",
+        "\n",
+        "# Training constants.\n",
+        "TRAINING_JOB_PREFIX = \"train\"\n",
+        "TRAIN_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-train\"\n",
+        "TRAIN_MACHINE_TYPE = \"n1-highmem-32\"\n",
+        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAIN_NUM_GPU = 8\n",
+        "\n",
+        "# Evaluation constants.\n",
+        "EVALUATION_METRIC = \"accuracy\"\n",
+        "\n",
+        "# Export constants.\n",
+        "EXPORT_JOB_PREFIX = \"export\"\n",
+        "EXPORT_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-model-export\"\n",
+        "EXPORT_MACHINE_TYPE = \"n1-highmem-8\"\n",
+        "\n",
+        "# Prediction constants.\n",
+        "# You can adjust accelerator types and machine types to get faster predictions.\n",
+        "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-serve\"\n",
+        "PREDICTION_PORT = 8501\n",
+        "PREDICTION_ACCELERATOR_COUNT = 1\n",
+        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
+        "PREDICTION_JOB_PREFIX = \"predict\""
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZZFPe_GezXg8"
+      },
+      "source": [
+        "### Define common helper functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "from datetime import datetime\n",
+        "from typing import Any\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "import yaml\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Returns a timestamped job name with the given prefix.\"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def get_label_map(label_map_yaml_filepath: str) -> tuple[dict[int, str], int]:\n",
+        "    \"\"\"Reads label map from a YAML file and returns the label map with the number of classes.\"\"\"\n",
+        "    with tf.io.gfile.GFile(label_map_yaml_filepath, \"rb\") as input_file:\n",
+        "        label_map = yaml.safe_load(input_file.read())[\"label_map\"]\n",
+        "    num_classes = max(label_map.keys()) + 1\n",
+        "    return label_map, num_classes\n",
+        "\n",
+        "\n",
+        "def get_best_trial(\n",
+        "    model_di: str, max_trial_count: int, evaluation_metric: str\n",
+        ") -> tuple[str, Any]:\n",
+        "    \"\"\"Finds the best trial directory and eval results from a hyperparameter tuning job.\"\"\"\n",
+        "    best_trial_dir = \"\"\n",
+        "    best_trial_evaluation_results = {}\n",
+        "    best_performance = -1\n",
+        "\n",
+        "    for i in range(max_trial_count):\n",
+        "        current_trial = i + 1\n",
+        "        current_trial_dir = os.path.join(model_dir, \"trial_\" + str(current_trial))\n",
+        "        current_trial_best_ckpt_dir = os.path.join(current_trial_dir, \"best_ckpt\")\n",
+        "        current_trial_best_ckpt_evaluation_filepath = os.path.join(\n",
+        "            current_trial_best_ckpt_dir, \"info.json\"\n",
+        "        )\n",
+        "        with tf.io.gfile.GFile(current_trial_best_ckpt_evaluation_filepath, \"rb\") as f:\n",
+        "            eval_metric_results = json.load(f)\n",
+        "            current_performance = eval_metric_results[evaluation_metric]\n",
+        "            if current_performance > best_performance:\n",
+        "                best_performance = current_performance\n",
+        "                best_trial_dir = current_trial_dir\n",
+        "                best_trial_evaluation_results = eval_metric_results\n",
+        "    return best_trial_dir, best_trial_evaluation_results\n",
+        "\n",
+        "\n",
+        "def print_response_instance(json_str: str, label_map: dict[int, str]):\n",
+        "    \"\"\"Prints summary of a prediction JSON result from the model response.\"\"\"\n",
+        "    json_obj = json.loads(json_str)\n",
+        "    if \"prediction\" not in json_obj:\n",
+        "        print(\"Error:\", json_str)\n",
+        "        return\n",
+        "    instance = json_obj[\"instance\"]\n",
+        "    prediction = json_obj[\"prediction\"]\n",
+        "    gcs_uri = instance[\"content\"]\n",
+        "    time_start = instance.get(\"timeSegmentStart\", \"0.0s\")\n",
+        "    time_end = instance.get(\"timeSegmentEnd\", \"Infinity\")\n",
+        "    print(f\"---------- Predict {gcs_uri}, {time_start} to {time_end}:\")\n",
+        "    for predicted in prediction:\n",
+        "        time = predicted[\"timeSegmentStart\"]\n",
+        "        label = label_map[predicted[\"label\"]]\n",
+        "        confidence = predicted[\"confidence\"]\n",
+        "        print(f\"At {time}, detected {label} with {confidence} confidence.\")\n",
+        "\n",
+        "\n",
+        "def find_checkpoint_in_dir(checkpoint_dir: str) -> str:\n",
+        "    \"\"\"Finds a checkpoint path relative to the directory.\"\"\"\n",
+        "    for root, dirs, files in tf.io.gfile.walk(checkpoint_dir):\n",
+        "        for file in files:\n",
+        "            if file.endswith(\".index\"):\n",
+        "                return os.path.join(root, os.path.splitext(file)[0])\n",
+        "\n",
+        "\n",
+        "def upload_checkpoint_to_gcs(checkpoint_url: str) -> str:\n",
+        "    \"\"\"Uploads a compressed .tar.gz checkpoint at the given URL to Cloud Storage.\"\"\"\n",
+        "    filename = os.path.basename(checkpoint_url)\n",
+        "    checkpoint_name = filename.replace(\".tar.gz\", \"\")\n",
+        "    print(\"Download checkpoint from\", checkpoint_url, \"and store to\", CHECKPOINT_BUCKET)\n",
+        "    ! wget $checkpoint_url -O $filename\n",
+        "    ! mkdir -p $checkpoint_name\n",
+        "    ! tar -xvzf $filename -C $checkpoint_name\n",
+        "\n",
+        "    checkpoint_path = find_checkpoint_in_dir(checkpoint_name)\n",
+        "    checkpoint_path = os.path.relpath(checkpoint_path, checkpoint_name)\n",
+        "\n",
+        "    ! gsutil cp -r $checkpoint_name $CHECKPOINT_BUCKET/\n",
+        "    checkpoint_uri = os.path.join(CHECKPOINT_BUCKET, checkpoint_name, checkpoint_path)\n",
+        "    print(\"Checkpoint uploaded to\", checkpoint_uri)\n",
+        "    return checkpoint_uri\n",
+        "\n",
+        "\n",
+        "def upload_config_to_gcs(url: str) -> str:\n",
+        "    \"\"\"Uploads a config file at the given URL to Cloud Storage.\"\"\"\n",
+        "    filename = os.path.basename(url)\n",
+        "    destination = os.path.join(CONFIG_DIR, filename)\n",
+        "    print(\"Copy\", url, \"to\", destination)\n",
+        "    ! wget \"$url\" -O \"$filename\"\n",
+        "    ! gsutil cp \"$filename\" \"$destination\"\n",
+        "    return destination"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RB_xY9ipr7ZU"
+      },
+      "source": [
+        "## Train new models\n",
+        "This section shows how to train new models.\n",
+        "1. Convert input data to training formats\n",
+        "2. Create hyperparameter tuning jobs to train new models\n",
+        "3. Find and export best models\n",
+        "\n",
+        "If you already trained models, please go to the section `Test Trained models`.\n",
+        "\n",
+        "Please select a model:\n",
+        "* `model_id`: MoViNet model variant ID, one of `a0`, `a1`, `a2`, `a3`, `a4`, `a5`. The model with a larger number requires more resources to train, and is expected to have a higher accuracy and latency. Here, we use `a3` for demonstration purpose. **`a0`, `a1`, and `a2` are not recommended for now as we are currently investigating some inference issues with them.**\n",
+        "* `model_mode`: MoViNet model type, either `base` or `stream`. The base model has a slightly higher accuracy, while the streaming model is optimized for streaming and faster CPU inference. See [official MoViNet docs](https://github.com/tensorflow/models/tree/master/official/projects/movinet) for more information.\n",
+        "\n",
+        "**Note**: The prediction container only supports base model (non-streaming) for now. If you train a streaming model, you need to download the model and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3Ry1mw6AHLTy"
+      },
+      "outputs": [],
+      "source": [
+        "model_id = \"a3\"  # @param [\"a0\", \"a1\", \"a2\", \"a3\", \"a4\", \"a5\"]\n",
+        "model_mode = \"base\"  # @param [\"base\", \"stream\"]\n",
+        "is_stream = model_mode == \"stream\"\n",
+        "model_name = f\"movinet_{model_id}_{model_mode}\"\n",
+        "image_size = IMAGE_SIZES[model_id]\n",
+        "\n",
+        "if is_stream:\n",
+        "    export_container_args = {\n",
+        "        \"conv_type\": \"2plus1d\",\n",
+        "        \"se_type\": \"2plus3d\",\n",
+        "        \"activation\": \"hard_swish\",\n",
+        "        \"gating_activation\": \"hard_sigmoid\",\n",
+        "        \"use_positional_encoding\": model_id in {\"a3\", \"a4\", \"a5\"},\n",
+        "    }\n",
+        "else:\n",
+        "    export_container_args = {\n",
+        "        \"conv_type\": \"3d\",\n",
+        "        \"se_type\": \"3d\",\n",
+        "        \"activation\": \"swish\",\n",
+        "        \"gating_activation\": \"sigmoid\",\n",
+        "        \"use_positional_encoding\": False,\n",
+        "    }"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Prepare data in the format as described [here](https://cloud.google.com/vertex-ai/docs/video-data/action-recognition/prepare-data), and then convert them to the training formats by running the cell below:\n",
+        "\n",
+        "* `input_file_path`: The input file path to the prepared data.\n",
+        "* `input_file_type`: The input file type, such as `csv` or `jsonl`.\n",
+        "* `output_fps`: The sampling rate of the video; Frames per second.\n",
+        "* `num_frames`: Number of frame to sample around keyframe inputs.\n",
+        "* `min_duration_sec`: Minimum duration in seconds for sampling video clips around keyframe inputs. This is for validation purpose - an error will be thrown if there is not enough context around a keyframe.\n",
+        "* `pos_neg_ratio`: Sampling ratio between positive and negative segments. For example, a pos_neg_ratio of 0.5 samples 1 negative instance every 2 positive instances.\n",
+        "* `split_ratio`: Three comma separated floats indicating the proportion of data to split into train/validation/test. They must add up to 1.\n",
+        "* `num_shard`: Three comma separated integers indicating the shards for train/validation/test.\n",
+        "\n",
+        "**Note**: For JSONL input, please use `aiplatform.googleapis.com/ml_use` instead of `ml_use` as the JSON key for ML use in `dataItemResourceLabels`. This is to be consistent with other objectives."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "# This job will convert input data as training format, with given split ratios\n",
+        "# and number of shards on train/test/validation.\n",
+        "\n",
+        "data_converter_job_name = get_job_name_with_datetime(\n",
+        "    DATA_CONVERTER_JOB_PREFIX + \"_\" + OBJECTIVE\n",
+        ")\n",
+        "\n",
+        "input_file_path = \"\"  # @param {type:\"string\"}\n",
+        "input_file_type = \"csv\"  # @param [\"csv\", \"jsonl\"]\n",
+        "output_fps = 10  # @param {type:\"integer\"}\n",
+        "num_frames = 32  # @param {type:\"integer\"}\n",
+        "min_duration_sec = 1.0  # @param {type:\"number\"}\n",
+        "pos_neg_ratio = 1.0  # @param {type:\"number\"}\n",
+        "split_ratio = \"0.8,0.1,0.1\"\n",
+        "num_shard = \"10,10,10\"\n",
+        "data_converter_output_dir = os.path.join(BUCKET_URI, data_converter_job_name)\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": DATA_CONVERTER_MACHINE_TYPE,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": DATA_CONVERTER_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--input_file_path=%s\" % input_file_path,\n",
+        "                \"--input_file_type=%s\" % input_file_type,\n",
+        "                \"--objective=%s\" % OBJECTIVE,\n",
+        "                \"--num_shard=%s\" % num_shard,\n",
+        "                \"--split_ratio=%s\" % split_ratio,\n",
+        "                \"--output_dir=%s\" % data_converter_output_dir,\n",
+        "                \"--output_fps=%d\" % output_fps,\n",
+        "                \"--num_frames=%d\" % num_frames,\n",
+        "                \"--min_duration_sec=%f\" % min_duration_sec,\n",
+        "                \"--pos_neg_ratio=%f\" % pos_neg_ratio,\n",
+        "                \"--output_shape=%d,%d\" % (image_size, image_size),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "data_converter_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=data_converter_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "data_converter_custom_job.run()\n",
+        "\n",
+        "input_train_data_path = os.path.join(data_converter_output_dir, \"train.tfrecord*\")\n",
+        "input_validation_data_path = os.path.join(data_converter_output_dir, \"val.tfrecord*\")\n",
+        "label_map_path = os.path.join(data_converter_output_dir, \"label_map.yaml\")\n",
+        "print(\"input_train_data_path for training: \", input_train_data_path)\n",
+        "print(\"input_validation_data_path for training: \", input_validation_data_path)\n",
+        "print(\"label_map_path for prediction: \", label_map_path)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Create a Vertex AI custom job with hyperparameter tuning\n",
+        "\n",
+        "You use the Vertex AI SDK to create and run the [hyperparameter tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) with Vertex AI Model Garden training docker images.\n",
+        "\n",
+        "#### Define the following specifications\n",
+        "\n",
+        "* `worker_pool_specs`: A list of dictionaries specifying the machine type and docker image. This example defines a single node cluster with one `n1-highmem-32` machine with 8 `NVIDIA_TESLA_V100` GPUs.\n",
+        "\n",
+        "  **Note**: We recommend using 8 GPUs for MoViNet-A2 and larger. Since loading video data requires a lot of GPU memory, it is recommended to experiment with a small batch size first.\n",
+        "* `parameter_spec`: Dictionary specifying the parameters to optimize. The dictionary key is the string assigned to the command line argument for each hyperparameter in your training application code, and the dictionary value is the parameter specification. The parameter specification includes the type, min/max values, and scale for the hyperparameter.\n",
+        "* `metric_spec`: Dictionary specifying the metric to optimize. The dictionary key is the `hyperparameter_metric_tag` that you set in your training application code, and the value is the optimization goal."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
+        "\n",
+        "# Input train and validation datasets can be found from the section above\n",
+        "# `Prepare input data for training`.\n",
+        "# Or, set prepared datasets paths if already exist.\n",
+        "# input_train_data_path = \"\"\n",
+        "# input_validation_data_path = \"\"\n",
+        "# label_map_path = \"\"\n",
+        "\n",
+        "train_job_name = get_job_name_with_datetime(f\"{TRAINING_JOB_PREFIX}_{model_name}\")\n",
+        "model_dir = os.path.join(BUCKET_URI, train_job_name)\n",
+        "label_map, num_classes = get_label_map(label_map_path)\n",
+        "\n",
+        "# Uploads pretained checkpoint to GCS bucket.\n",
+        "init_checkpoint = f\"https://storage.googleapis.com/tf_model_garden/vision/movinet/{model_name}_with_backbone.tar.gz\"\n",
+        "init_checkpoint = upload_checkpoint_to_gcs(init_checkpoint)\n",
+        "\n",
+        "# Uploads config file according to model_id and streaming options.\n",
+        "config_file = f\"{model_id}_stream\" if is_stream else model_id\n",
+        "config_file = f\"https://raw.githubusercontent.com/tensorflow/models/master/official/projects/movinet/configs/yaml/movinet_{config_file}_gpu.yaml\"\n",
+        "config_file = upload_config_to_gcs(config_file)\n",
+        "\n",
+        "# The parameters here are mainly for demonstration purpose. Please update them\n",
+        "# for better performance.\n",
+        "trainer_args = {\n",
+        "    \"experiment\": \"movinet_kinetics600\",\n",
+        "    \"config_file\": config_file,\n",
+        "    \"input_train_data_path\": input_train_data_path,\n",
+        "    \"input_validation_data_path\": input_validation_data_path,\n",
+        "    \"init_checkpoint\": init_checkpoint,\n",
+        "    \"model_dir\": model_dir,\n",
+        "    \"num_classes\": num_classes,\n",
+        "    \"global_batch_size\": 16,\n",
+        "    \"prefetch_buffer_size\": 16,\n",
+        "    \"shuffle_buffer_size\": 32,\n",
+        "    \"train_steps\": 2000,\n",
+        "}\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
+        "            # Each training job uses TRAIN_NUM_GPU GPUs.\n",
+        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
+        "            \"args\": [\n",
+        "                \"--mode=train_and_eval\",\n",
+        "                \"--params_override=runtime.num_gpus=%d\" % TRAIN_NUM_GPU,\n",
+        "            ]\n",
+        "            + [\"--{}={}\".format(k, v) for k, v in trainer_args.items()],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "metric_spec = {\"model_performance\": \"maximize\"}\n",
+        "\n",
+        "# These learning rates might not be optimal for your selected model type; To\n",
+        "# tune learning rates, try hpt.DoubleParameterSpec with more trials.\n",
+        "LEARNING_RATES = [1e-3, 3e-3]\n",
+        "MAX_TRIAL_COUNT = len(LEARNING_RATES)\n",
+        "parameter_spec = {\n",
+        "    \"learning_rate\": hpt.DiscreteParameterSpec(values=LEARNING_RATES, scale=\"linear\"),\n",
+        "}\n",
+        "\n",
+        "print(worker_pool_specs, metric_spec, parameter_spec)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "#### Run the hyperparameter tuning job\n",
+        "* `max_trial_count`: Sets an upper bound on the number of trials the service will run. The recommended practice is to start with a smaller number of trials and get a sense of how impactful your chosen hyperparameters are before scaling up.\n",
+        "\n",
+        "* `parallel_trial_count`:  If you use parallel trials, the service provisions multiple training processing clusters. The worker pool spec that you specify when creating the job is used for each individual training cluster.  Increasing the number of parallel trials reduces the amount of time the hyperparameter tuning job takes to run; however, it can reduce the effectiveness of the job overall. This is because the default tuning strategy uses results of previous trials to inform the assignment of values in subsequent trials.\n",
+        "\n",
+        "* `search_algorithm`: The available search algorithms are grid, random, or default (None). The default option applies Bayesian optimization to search the space of possible hyperparameter values and is the recommended algorithm.\n",
+        "\n",
+        "Click on the generated link in the output to see your run in the Cloud Console."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "train_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=train_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "train_hpt_job = aiplatform.HyperparameterTuningJob(\n",
+        "    display_name=train_job_name,\n",
+        "    custom_job=train_custom_job,\n",
+        "    metric_spec=metric_spec,\n",
+        "    parameter_spec=parameter_spec,\n",
+        "    max_trial_count=MAX_TRIAL_COUNT,\n",
+        "    parallel_trial_count=1,\n",
+        "    project=PROJECT_ID,\n",
+        "    search_algorithm=None,\n",
+        ")\n",
+        "\n",
+        "train_hpt_job.run()\n",
+        "\n",
+        "print(\"model_dir is:\", model_dir)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vugUfJEC2HrK"
+      },
+      "source": [
+        "### Export model in Tensorflow SavedModel format"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "09Rz1AYspK19"
+      },
+      "outputs": [],
+      "source": [
+        "# This job will export models from TF checkpoints to TF saved model format.\n",
+        "# model_dir is from the section above.\n",
+        "best_trial_dir, best_trial_evaluation_results = get_best_trial(\n",
+        "    model_dir, MAX_TRIAL_COUNT, EVALUATION_METRIC\n",
+        ")\n",
+        "best_checkpoint_path = find_checkpoint_in_dir(f\"{best_trial_dir}/best_ckpt/\")\n",
+        "print(\"best_trial_dir: \", best_trial_dir)\n",
+        "print(\"best_trial_evaluation_results: \", best_trial_evaluation_results)\n",
+        "print(\"best_checkpoint: \", best_checkpoint_path)\n",
+        "\n",
+        "container_args = {\n",
+        "    \"export_path\": f\"{model_dir}/best_model\",\n",
+        "    \"model_id\": model_id,\n",
+        "    \"num_classes\": num_classes,\n",
+        "    \"causal\": is_stream,\n",
+        "    \"checkpoint_path\": best_checkpoint_path,\n",
+        "    \"assert_checkpoint_objects_matched\": False,\n",
+        "    **export_container_args,\n",
+        "}\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": EXPORT_MACHINE_TYPE,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": EXPORT_CONTAINER_URI,\n",
+        "            \"args\": [\"--{}={}\".format(k, v) for k, v in container_args.items()],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "model_export_job_name = get_job_name_with_datetime(EXPORT_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
+        "model_export_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=model_export_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "model_export_custom_job.run()\n",
+        "\n",
+        "print(\"best model is saved to: \", container_args[\"export_path\"])"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "## Test trained models\n",
+        "This section shows the way to test with trained models.\n",
+        "1. Upload and deploy models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
+        "2. Run batch predictions\n",
+        "\n",
+        "**Note:** The prediction container only works with the base model. If you trained a streaming model, download the model from the exported path and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gdlca3BOypXU"
+      },
+      "source": [
+        "### Upload model to Vertex AI Model Registry\n",
+        "\n",
+        "The following cell uploads the trained model to Vertex AI Model Registry. Skip it if you want to run batch predictions on an already uploaded model instead.\n",
+        "\n",
+        "#### Configurable environment variables\n",
+        "\n",
+        "* `MODEL_PATH`: Cloud Storage URI to the MoViNet model.\n",
+        "* `BATCH_SIZE`: Batch size for inference. Use a larger value to accelerate GPU prediction.\n",
+        "* `NUM_FRAMES`: Number of frames for a single prediction with the model.\n",
+        "* `FPS`: Video sampling frame per second.\n",
+        "* `OVERLAP_FRAMES`: Allowed overlapping frames between consecutive prediction windows. Set a smaller value for faster inference but less accurate."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "serving_env = {\n",
+        "    \"MODEL_PATH\": container_args[\"export_path\"],\n",
+        "    \"BATCH_SIZE\": 1,\n",
+        "    \"NUM_FRAMES\": 32,\n",
+        "    \"FPS\": output_fps,\n",
+        "    \"OVERLAP_FRAMES\": 24,\n",
+        "    \"OBJECTIVE\": OBJECTIVE,\n",
+        "    \"IMAGE_WIDTH\": image_size,\n",
+        "    \"IMAGE_HEIGHT\": image_size,\n",
+        "}\n",
+        "\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=model_name,\n",
+        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
+        "    serving_container_ports=[PREDICTION_PORT],\n",
+        "    serving_container_predict_route=\"/predict\",\n",
+        "    serving_container_health_route=\"/ping\",\n",
+        "    serving_container_environment_variables=serving_env,\n",
+        ")\n",
+        "\n",
+        "model.wait()\n",
+        "\n",
+        "print(\"The uploaded model name is: \", model_name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2b47e629a01"
+      },
+      "source": [
+        "Alternatively, uncomment the following cell to use an already uploaded model. Replace the model name string with that of the existing model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aa503a565b8f"
+      },
+      "outputs": [],
+      "source": [
+        "# model = aiplatform.Model(\"projects/123456789/locations/us-central1/models/12345678901234567890\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9SZsKGeS3x6S"
+      },
+      "source": [
+        "### Run batch predictions\n",
+        "\n",
+        "We will now run batch predictions with the trained MoViNet action recognition model with [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions).\n",
+        "\n",
+        "Please prepare an input JSONL file where each line follows [this format](https://cloud.google.com/vertex-ai/docs/video-data/action-recognition/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket. The service account should have read access to the buckets containing the trained model and the input data. See [Service accounts overview](https://cloud.google.com/iam/docs/service-account-overview) for more information.\n",
+        "\n",
+        "The [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions) has a default timeout of 10 minutes. Therefore, please make sure the input video clip is around 5 minutes at 5~10 FPS or you may experience a timeout error. To use this model at a larger scale beyond this notebook demontration, you can try one of the following:\n",
+        "\n",
+        "- Pull the serving docker image to a VM or a local machine and send prediction requests directly.\n",
+        "- To process more data concurrently, write a custom [DataFlow](https://cloud.google.com/dataflow) pipeline to send prediction requests to the movinet serving container.\n",
+        "- Divide videos into 5-minute clips and run batch prediction with a small batch size."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vbIW9me1F2RY"
+      },
+      "outputs": [],
+      "source": [
+        "# Path to the prediction input JSONL file.\n",
+        "test_jsonl_path = \"\"  # @param {type:\"string\"}\n",
+        "# Full service account name with the suffix `gserviceaccount.com`.\n",
+        "batch_predict_service_account = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "predict_job_name = get_job_name_with_datetime(f\"{PREDICTION_JOB_PREFIX}_{model_name}\")\n",
+        "predict_destination_prefix = os.path.join(STAGING_BUCKET, predict_job_name)\n",
+        "\n",
+        "batch_prediction_job = model.batch_predict(\n",
+        "    job_display_name=predict_job_name,\n",
+        "    gcs_source=test_jsonl_path,\n",
+        "    gcs_destination_prefix=predict_destination_prefix,\n",
+        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    accelerator_count=PREDICTION_ACCELERATOR_COUNT,\n",
+        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    max_replica_count=1,\n",
+        "    service_account=batch_predict_service_account,\n",
+        ")\n",
+        "\n",
+        "batch_prediction_job.wait()\n",
+        "\n",
+        "print(batch_prediction_job.display_name)\n",
+        "print(batch_prediction_job.resource_name)\n",
+        "print(batch_prediction_job.state)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ik-XPjfx9OCE"
+      },
+      "source": [
+        "You can then read the prediction response JSONL files in the output directory:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tdkW9e5B9OU1"
+      },
+      "outputs": [],
+      "source": [
+        "# The label map file was generated from the section above (`Prepare input data for training`).\n",
+        "for file in tf.io.gfile.glob(os.path.join(predict_destination_prefix, \"*/*\")):\n",
+        "    with tf.io.gfile.GFile(file, \"r\") as f:\n",
+        "        for line in f:\n",
+        "            print_response_instance(line, label_map)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete the trained model.\n",
+        "model.delete()\n",
+        "# Delete custom and hpt jobs.\n",
+        "if data_converter_custom_job.list(filter=f'display_name=\"{data_converter_job_name}\"'):\n",
+        "    data_converter_custom_job.delete()\n",
+        "if train_hpt_job.list(filter=f'display_name=\"{train_job_name}\"'):\n",
+        "    train_hpt_job.delete()\n",
+        "if model_export_custom_job.list(filter=f'display_name=\"{model_export_job_name}\"'):\n",
+        "    model_export_custom_job.delete()\n",
+        "if batch_prediction_job.list(filter=f'display_name=\"{predict_job_name}\"'):\n",
+        "    batch_prediction_job.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_movinet_action_recognition.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb

@@ -0,0 +1,912 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden MoViNet video clip classification\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>                                                                                               <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [MoViNet](https://github.com/tensorflow/models/tree/master/official/projects/movinet) in Vertex AI Model Garden.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Train new models\n",
+        "  * Convert input data to training formats\n",
+        "  * Create [hyperparameter tuning jobs](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to train new models\n",
+        "  * Find and export best models\n",
+        "\n",
+        "* Test trained models\n",
+        "  * Upload models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
+        "  * Run batch predictions\n",
+        "\n",
+        "* Clean up resources\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "### Colab Only\n",
+        "Run the following commands for Colab or skip this section if you use Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for running batch predictions with the fine tuned model.\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9wExiMUxFk91"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# The GCP project ID for experiments.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Bucket URI with gs:// prefix.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# You can choose a region from https://cloud.google.com/about/locations.\n",
+        "# Only regions prefixed by \"us\", \"asia\", or \"europe\" are supported.\n",
+        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
+        "assert REGION_PREFIX in (\n",
+        "    \"us\",\n",
+        "    \"europe\",\n",
+        "    \"asia\",\n",
+        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "CHECKPOINT_BUCKET = os.path.join(BUCKET_URI, \"ckpt\")\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Download config files.\n",
+        "CONFIG_DIR = os.path.join(BUCKET_URI, \"config\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "n6IFz75WGCam"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "OBJECTIVE = \"vcn\"\n",
+        "\n",
+        "# Data converter constants.\n",
+        "DATA_CONVERTER_JOB_PREFIX = \"data_converter\"\n",
+        "DATA_CONVERTER_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/data-converter\"\n",
+        "DATA_CONVERTER_MACHINE_TYPE = \"n1-highmem-8\"\n",
+        "\n",
+        "# Training constants.\n",
+        "TRAINING_JOB_PREFIX = \"train\"\n",
+        "TRAIN_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-train\"\n",
+        "TRAIN_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAIN_NUM_GPU = 2\n",
+        "\n",
+        "# Evaluation constants.\n",
+        "EVALUATION_METRIC = \"accuracy\"\n",
+        "\n",
+        "# Export constants.\n",
+        "EXPORT_JOB_PREFIX = \"export\"\n",
+        "EXPORT_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-model-export\"\n",
+        "EXPORT_MACHINE_TYPE = \"n1-highmem-8\"\n",
+        "\n",
+        "# Prediction constants.\n",
+        "# You can adjust accelerator types and machine types to get faster predictions.\n",
+        "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-serve\"\n",
+        "PREDICTION_PORT = 8501\n",
+        "PREDICTION_ACCELERATOR_COUNT = 1\n",
+        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
+        "PREDICTION_JOB_PREFIX = \"predict\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZZFPe_GezXg8"
+      },
+      "source": [
+        "### Define common helper functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "from datetime import datetime\n",
+        "from typing import Any\n",
+        "\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "import yaml\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Returns a timestamped job name with the given prefix.\"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def print_response_instance(json_str: str, label_map: dict[int, str]):\n",
+        "    \"\"\"Prints summary of a prediction JSON result from the model response.\"\"\"\n",
+        "    json_obj = json.loads(json_str)\n",
+        "    if \"prediction\" not in json_obj:\n",
+        "        print(\"Error:\", json_str)\n",
+        "        return\n",
+        "    instance = json_obj[\"instance\"]\n",
+        "    prediction = json_obj[\"prediction\"]\n",
+        "    gcs_uri = instance[\"content\"]\n",
+        "    time_start = instance.get(\"timeSegmentStart\", \"0.0s\")\n",
+        "    time_end = instance.get(\"timeSegmentEnd\", \"Infinity\")\n",
+        "    max_idx = np.argmax(prediction)\n",
+        "    print(f\"{gcs_uri} {time_start}-{time_end}:\", label_map[max_idx])\n",
+        "\n",
+        "\n",
+        "def get_label_map(label_map_yaml_filepath: str) -> tuple[dict[int, str], int]:\n",
+        "    \"\"\"Reads label map from a YAML file and returns the label map with the number of classes.\"\"\"\n",
+        "    with tf.io.gfile.GFile(label_map_yaml_filepath, \"rb\") as input_file:\n",
+        "        label_map = yaml.safe_load(input_file.read())[\"label_map\"]\n",
+        "    num_classes = max(label_map.keys()) + 1\n",
+        "    return label_map, num_classes\n",
+        "\n",
+        "\n",
+        "def get_best_trial(\n",
+        "    model_di: str, max_trial_count: int, evaluation_metric: str\n",
+        ") -> tuple[str, Any]:\n",
+        "    \"\"\"Finds the best trial directory and eval results from a hyperparameter tuning job.\"\"\"\n",
+        "    best_trial_dir = \"\"\n",
+        "    best_trial_evaluation_results = {}\n",
+        "    best_performance = -1\n",
+        "\n",
+        "    for i in range(max_trial_count):\n",
+        "        current_trial = i + 1\n",
+        "        current_trial_dir = os.path.join(model_dir, \"trial_\" + str(current_trial))\n",
+        "        current_trial_best_ckpt_dir = os.path.join(current_trial_dir, \"best_ckpt\")\n",
+        "        current_trial_best_ckpt_evaluation_filepath = os.path.join(\n",
+        "            current_trial_best_ckpt_dir, \"info.json\"\n",
+        "        )\n",
+        "        with tf.io.gfile.GFile(current_trial_best_ckpt_evaluation_filepath, \"rb\") as f:\n",
+        "            eval_metric_results = json.load(f)\n",
+        "            current_performance = eval_metric_results[evaluation_metric]\n",
+        "            if current_performance > best_performance:\n",
+        "                best_performance = current_performance\n",
+        "                best_trial_dir = current_trial_dir\n",
+        "                best_trial_evaluation_results = eval_metric_results\n",
+        "    return best_trial_dir, best_trial_evaluation_results\n",
+        "\n",
+        "\n",
+        "def find_checkpoint_in_dir(checkpoint_dir: str) -> str:\n",
+        "    \"\"\"Finds a checkpoint path relative to the directory.\"\"\"\n",
+        "    for root, dirs, files in tf.io.gfile.walk(checkpoint_dir):\n",
+        "        for file in files:\n",
+        "            if file.endswith(\".index\"):\n",
+        "                return os.path.join(root, os.path.splitext(file)[0])\n",
+        "\n",
+        "\n",
+        "def upload_checkpoint_to_gcs(checkpoint_url: str) -> str:\n",
+        "    \"\"\"Uploads a compressed .tar.gz checkpoint at the given URL to Cloud Storage.\"\"\"\n",
+        "    filename = os.path.basename(checkpoint_url)\n",
+        "    checkpoint_name = filename.replace(\".tar.gz\", \"\")\n",
+        "    print(\"Download checkpoint from\", checkpoint_url, \"and store to\", CHECKPOINT_BUCKET)\n",
+        "    ! wget $checkpoint_url -O $filename\n",
+        "    ! mkdir -p $checkpoint_name\n",
+        "    ! tar -xvzf $filename -C $checkpoint_name\n",
+        "\n",
+        "    checkpoint_path = find_checkpoint_in_dir(checkpoint_name)\n",
+        "    checkpoint_path = os.path.relpath(checkpoint_path, checkpoint_name)\n",
+        "\n",
+        "    ! gsutil cp -r $checkpoint_name $CHECKPOINT_BUCKET/\n",
+        "    checkpoint_uri = os.path.join(CHECKPOINT_BUCKET, checkpoint_name, checkpoint_path)\n",
+        "    print(\"Checkpoint uploaded to\", checkpoint_uri)\n",
+        "    return checkpoint_uri\n",
+        "\n",
+        "\n",
+        "def upload_config_to_gcs(url: str) -> str:\n",
+        "    \"\"\"Uploads a config file at the given URL to Cloud Storage.\"\"\"\n",
+        "    filename = os.path.basename(url)\n",
+        "    destination = os.path.join(CONFIG_DIR, filename)\n",
+        "    print(\"Copy\", url, \"to\", destination)\n",
+        "    ! wget \"$url\" -O \"$filename\"\n",
+        "    ! gsutil cp \"$filename\" \"$destination\"\n",
+        "    return destination"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RB_xY9ipr7ZU"
+      },
+      "source": [
+        "## Train new models\n",
+        "This section shows how to train new models.\n",
+        "1. Convert input data to training formats\n",
+        "2. Create hyperparameter tuning jobs to train new models\n",
+        "3. Find and export best models\n",
+        "\n",
+        "If you already trained models, please go to the section `Test Trained models`.\n",
+        "\n",
+        "Please select a model:\n",
+        "* `model_id`: MoViNet model variant ID, one of `a0`, `a1`, `a2`, `a3`, `a4`, `a5`. The model with a larger number requires more resources to train, and is expected to have a higher accuracy and latency. Here, we use `a0` for demonstration purpose.\n",
+        "* `model_mode`: MoViNet model type, either `base` or `stream`. The base model has a slightly higher accuracy, while the streaming model is optimized for streaming and faster CPU inference. See [official MoViNet docs](https://github.com/tensorflow/models/tree/master/official/projects/movinet) for more information.\n",
+        "\n",
+        "**Note**: The prediction container only supports base model (non-streaming) for now. If you train a streaming model, you need to download the model and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3Ry1mw6AHLTy"
+      },
+      "outputs": [],
+      "source": [
+        "model_id = \"a0\"  # @param [\"a0\", \"a1\", \"a2\", \"a3\", \"a4\", \"a5\"]\n",
+        "model_mode = \"base\"  # @param [\"base\", \"stream\"]\n",
+        "is_stream = model_mode == \"stream\"\n",
+        "model_name = f\"movinet_{model_id}_{model_mode}\"\n",
+        "\n",
+        "if is_stream:\n",
+        "    export_container_args = {\n",
+        "        \"conv_type\": \"2plus1d\",\n",
+        "        \"se_type\": \"2plus3d\",\n",
+        "        \"activation\": \"hard_swish\",\n",
+        "        \"gating_activation\": \"hard_sigmoid\",\n",
+        "        \"use_positional_encoding\": model_id in {\"a3\", \"a4\", \"a5\"},\n",
+        "    }\n",
+        "else:\n",
+        "    export_container_args = {\n",
+        "        \"conv_type\": \"3d\",\n",
+        "        \"se_type\": \"3d\",\n",
+        "        \"activation\": \"swish\",\n",
+        "        \"gating_activation\": \"sigmoid\",\n",
+        "        \"use_positional_encoding\": False,\n",
+        "    }"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Prepare input data for training\n",
+        "\n",
+        "Prepare data in the format as described [here](https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data), and then convert them to the training formats by running the cell below:\n",
+        "\n",
+        "* `input_file_path`: The input file path to the prepared data.\n",
+        "* `input_file_type`: The input file type, such as `csv` or `jsonl`.\n",
+        "* `output_fps`: The sampling rate of the video; Frames per second.\n",
+        "* `split_ratio`: Three comma separated floats indicating the proportion of data to split into train/validation/test. They must add up to 1.\n",
+        "* `num_shard`: Three comma separated integers indicating the shards for train/validation/test."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "IndQ_m6ddUEM"
+      },
+      "outputs": [],
+      "source": [
+        "# This job will convert input data as training format, with given split ratios\n",
+        "# and number of shards on train/test/validation.\n",
+        "\n",
+        "data_converter_job_name = get_job_name_with_datetime(\n",
+        "    DATA_CONVERTER_JOB_PREFIX + \"_\" + OBJECTIVE\n",
+        ")\n",
+        "\n",
+        "input_file_path = \"\"  # @param {type:\"string\"}\n",
+        "input_file_type = \"csv\"  # @param [\"csv\", \"jsonl\"]\n",
+        "output_fps = 5  # @param {type:\"integer\"}\n",
+        "split_ratio = \"0.8,0.1,0.1\"\n",
+        "num_shard = \"10,10,10\"\n",
+        "data_converter_output_dir = os.path.join(BUCKET_URI, data_converter_job_name)\n",
+        "\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": DATA_CONVERTER_MACHINE_TYPE,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": DATA_CONVERTER_CONTAINER,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--input_file_path=%s\" % input_file_path,\n",
+        "                \"--input_file_type=%s\" % input_file_type,\n",
+        "                \"--objective=%s\" % OBJECTIVE,\n",
+        "                \"--num_shard=%s\" % num_shard,\n",
+        "                \"--split_ratio=%s\" % split_ratio,\n",
+        "                \"--output_dir=%s\" % data_converter_output_dir,\n",
+        "                \"--output_fps=%d\" % output_fps,\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "data_converter_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=data_converter_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "data_converter_custom_job.run()\n",
+        "\n",
+        "input_train_data_path = os.path.join(data_converter_output_dir, \"train.tfrecord*\")\n",
+        "input_validation_data_path = os.path.join(data_converter_output_dir, \"val.tfrecord*\")\n",
+        "label_map_path = os.path.join(data_converter_output_dir, \"label_map.yaml\")\n",
+        "print(\"input_train_data_path for training: \", input_train_data_path)\n",
+        "print(\"input_validation_data_path for training: \", input_validation_data_path)\n",
+        "print(\"label_map_path for prediction: \", label_map_path)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aaff6f5be7f6"
+      },
+      "source": [
+        "### Create a Vertex AI custom job with hyperparameter tuning\n",
+        "\n",
+        "You use the Vertex AI SDK to create and run the [hyperparameter tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) with Vertex AI Model Garden training docker images.\n",
+        "\n",
+        "#### Define the following specifications\n",
+        "\n",
+        "* `worker_pool_specs`: A list of dictionaries specifying the machine type and docker image. This example defines a single node cluster with one `n1-highmem-16` machine with 2 `NVIDIA_TESLA_V100` GPUs.\n",
+        "\n",
+        "  **Note**: We recommend using 8 GPUs for MoViNet-A2 and larger. Since loading video data requires a lot of GPU memory, it is recommended to experiment with a small batch size first.\n",
+        "* `parameter_spec`: Dictionary specifying the parameters to optimize. The dictionary key is the string assigned to the command line argument for each hyperparameter in your training application code, and the dictionary value is the parameter specification. The parameter specification includes the type, min/max values, and scale for the hyperparameter.\n",
+        "* `metric_spec`: Dictionary specifying the metric to optimize. The dictionary key is the `hyperparameter_metric_tag` that you set in your training application code, and the value is the optimization goal."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "um_XKbmpTaHx"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
+        "\n",
+        "# Input train and validation datasets can be found from the section above\n",
+        "# `Prepare input data for training`.\n",
+        "# Or, set prepared datasets paths if already exist.\n",
+        "# input_train_data_path = \"\"\n",
+        "# input_validation_data_path = \"\"\n",
+        "# label_map_path = \"\"\n",
+        "\n",
+        "train_job_name = get_job_name_with_datetime(f\"{TRAINING_JOB_PREFIX}_{model_name}\")\n",
+        "model_dir = os.path.join(BUCKET_URI, train_job_name)\n",
+        "label_map, num_classes = get_label_map(label_map_path)\n",
+        "\n",
+        "# Uploads pretained checkpoint to GCS bucket.\n",
+        "init_checkpoint = f\"https://storage.googleapis.com/tf_model_garden/vision/movinet/{model_name}_with_backbone.tar.gz\"\n",
+        "init_checkpoint = upload_checkpoint_to_gcs(init_checkpoint)\n",
+        "\n",
+        "# Uploads config file according to model_id and streaming options.\n",
+        "config_file = f\"{model_id}_stream\" if is_stream else model_id\n",
+        "config_file = f\"https://raw.githubusercontent.com/tensorflow/models/master/official/projects/movinet/configs/yaml/movinet_{config_file}_gpu.yaml\"\n",
+        "config_file = upload_config_to_gcs(config_file)\n",
+        "\n",
+        "# The parameters here are mainly for demonstration purpose. Please update them\n",
+        "# for better performance.\n",
+        "trainer_args = {\n",
+        "    \"experiment\": \"movinet_kinetics600\",\n",
+        "    \"config_file\": config_file,\n",
+        "    \"input_train_data_path\": input_train_data_path,\n",
+        "    \"input_validation_data_path\": input_validation_data_path,\n",
+        "    \"init_checkpoint\": init_checkpoint,\n",
+        "    \"model_dir\": model_dir,\n",
+        "    \"num_classes\": num_classes,\n",
+        "    \"global_batch_size\": 4,\n",
+        "    \"prefetch_buffer_size\": 8,\n",
+        "    \"shuffle_buffer_size\": 32,\n",
+        "    \"train_steps\": 2000,\n",
+        "}\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
+        "            # Each training job uses TRAIN_NUM_GPU GPUs.\n",
+        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
+        "            \"args\": [\n",
+        "                \"--mode=train_and_eval\",\n",
+        "                \"--params_override=runtime.num_gpus=%d\" % TRAIN_NUM_GPU,\n",
+        "            ]\n",
+        "            + [\"--{}={}\".format(k, v) for k, v in trainer_args.items()],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "metric_spec = {\"model_performance\": \"maximize\"}\n",
+        "\n",
+        "# These learning rates might not be optimal for your selected model type; To\n",
+        "# tune learning rates, try hpt.DoubleParameterSpec with more trials.\n",
+        "LEARNING_RATES = [1e-3, 3e-3]\n",
+        "MAX_TRIAL_COUNT = len(LEARNING_RATES)\n",
+        "parameter_spec = {\n",
+        "    \"learning_rate\": hpt.DiscreteParameterSpec(values=LEARNING_RATES, scale=\"linear\"),\n",
+        "}\n",
+        "\n",
+        "print(worker_pool_specs, metric_spec, parameter_spec)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HwcCjwlBTQIz"
+      },
+      "source": [
+        "#### Run the hyperparameter tuning job\n",
+        "* `max_trial_count`: Sets an upper bound on the number of trials the service will run. The recommended practice is to start with a smaller number of trials and get a sense of how impactful your chosen hyperparameters are before scaling up.\n",
+        "\n",
+        "* `parallel_trial_count`:  If you use parallel trials, the service provisions multiple training processing clusters. The worker pool spec that you specify when creating the job is used for each individual training cluster.  Increasing the number of parallel trials reduces the amount of time the hyperparameter tuning job takes to run; however, it can reduce the effectiveness of the job overall. This is because the default tuning strategy uses results of previous trials to inform the assignment of values in subsequent trials.\n",
+        "\n",
+        "* `search_algorithm`: The available search algorithms are grid, random, or default (None). The default option applies Bayesian optimization to search the space of possible hyperparameter values and is the recommended algorithm.\n",
+        "\n",
+        "Click on the generated link in the output to see your run in the Cloud Console."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "train_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=train_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "train_hpt_job = aiplatform.HyperparameterTuningJob(\n",
+        "    display_name=train_job_name,\n",
+        "    custom_job=train_custom_job,\n",
+        "    metric_spec=metric_spec,\n",
+        "    parameter_spec=parameter_spec,\n",
+        "    max_trial_count=MAX_TRIAL_COUNT,\n",
+        "    parallel_trial_count=MAX_TRIAL_COUNT,\n",
+        "    project=PROJECT_ID,\n",
+        "    search_algorithm=None,\n",
+        ")\n",
+        "\n",
+        "train_hpt_job.run()\n",
+        "\n",
+        "print(\"model_dir is:\", model_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vugUfJEC2HrK"
+      },
+      "source": [
+        "### Export model in Tensorflow SavedModel format"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "09Rz1AYspK19"
+      },
+      "outputs": [],
+      "source": [
+        "# This job will export models from TF checkpoints to TF saved model format.\n",
+        "# model_dir is from the section above.\n",
+        "best_trial_dir, best_trial_evaluation_results = get_best_trial(\n",
+        "    model_dir, MAX_TRIAL_COUNT, EVALUATION_METRIC\n",
+        ")\n",
+        "best_checkpoint_path = find_checkpoint_in_dir(f\"{best_trial_dir}/best_ckpt/\")\n",
+        "print(\"best_trial_dir: \", best_trial_dir)\n",
+        "print(\"best_trial_evaluation_results: \", best_trial_evaluation_results)\n",
+        "print(\"best_checkpoint: \", best_checkpoint_path)\n",
+        "\n",
+        "container_args = {\n",
+        "    \"export_path\": f\"{model_dir}/best_model\",\n",
+        "    \"model_id\": model_id,\n",
+        "    \"num_classes\": num_classes,\n",
+        "    \"causal\": is_stream,\n",
+        "    \"checkpoint_path\": best_checkpoint_path,\n",
+        "    \"assert_checkpoint_objects_matched\": False,\n",
+        "    **export_container_args,\n",
+        "}\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": EXPORT_MACHINE_TYPE,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": EXPORT_CONTAINER_URI,\n",
+        "            \"args\": [\"--{}={}\".format(k, v) for k, v in container_args.items()],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "model_export_job_name = get_job_name_with_datetime(EXPORT_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
+        "model_export_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=model_export_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "model_export_custom_job.run()\n",
+        "\n",
+        "print(\"best model is saved to: \", container_args[\"export_path\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g0BGaofgsMsy"
+      },
+      "source": [
+        "## Test trained models\n",
+        "This section shows the way to test with trained models.\n",
+        "1. Upload and deploy models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
+        "2. Run batch predictions\n",
+        "\n",
+        "**Note:** The prediction container only works with the base model. If you trained a streaming model, download the model from the exported path and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gdlca3BOypXU"
+      },
+      "source": [
+        "### Upload model to Vertex AI Model Registry\n",
+        "\n",
+        "The following cell uploads the trained model to Vertex AI Model Registry. Skip it if you want to run batch predictions on an already uploaded model instead.\n",
+        "\n",
+        "#### Configurable environment variables\n",
+        "\n",
+        "* `MODEL_PATH`: Cloud Storage URI to the MoViNet model.\n",
+        "* `BATCH_SIZE`: Batch size for inference. Use a larger value to accelerate GPU prediction.\n",
+        "* `NUM_FRAMES`: Number of frames for a single prediction with the model.\n",
+        "* `FPS`: Video sampling frame per second.\n",
+        "* `OVERLAP_FRAMES`: Allowed overlapping frames between consecutive prediction windows. Set a smaller value for faster inference but less accurate."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NYuQowyZEtxK"
+      },
+      "outputs": [],
+      "source": [
+        "serving_env = {\n",
+        "    \"MODEL_PATH\": container_args[\"export_path\"],\n",
+        "    \"BATCH_SIZE\": 1,  # Select a larger batch size to accelerate GPU prediction.\n",
+        "    \"NUM_FRAMES\": 32,\n",
+        "    \"FPS\": output_fps,\n",
+        "    \"OVERLAP_FRAMES\": 24,\n",
+        "    \"OBJECTIVE\": OBJECTIVE,\n",
+        "}\n",
+        "\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=model_name,\n",
+        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
+        "    serving_container_ports=[PREDICTION_PORT],\n",
+        "    serving_container_predict_route=\"/predict\",\n",
+        "    serving_container_health_route=\"/ping\",\n",
+        "    serving_container_environment_variables=serving_env,\n",
+        ")\n",
+        "\n",
+        "model.wait()\n",
+        "\n",
+        "print(\"The uploaded model name is: \", model_name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c2b47e629a01"
+      },
+      "source": [
+        "Alternatively, uncomment the following cell to use an already uploaded model. Replace the model name string with that of the existing model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aa503a565b8f"
+      },
+      "outputs": [],
+      "source": [
+        "# model = aiplatform.Model(\"projects/123456789/locations/us-central1/models/12345678901234567890\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9SZsKGeS3x6S"
+      },
+      "source": [
+        "### Run batch predictions\n",
+        "\n",
+        "We will now run batch predictions with the trained MoViNet clip classification model with [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions).\n",
+        "\n",
+        "Please prepare an input JSONL file where each line follows [this format](https://cloud.google.com/vertex-ai/docs/video-data/classification/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket. The service account should have read access to the buckets containing the trained model and the input data. See [Service accounts overview](https://cloud.google.com/iam/docs/service-account-overview) for more information."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vbIW9me1F2RY"
+      },
+      "outputs": [],
+      "source": [
+        "# Path to the prediction input JSONL file.\n",
+        "test_jsonl_path = \"\"  # @param {type:\"string\"}\n",
+        "# Full service account name with the suffix `gserviceaccount.com`.\n",
+        "batch_predict_service_account = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "predict_job_name = get_job_name_with_datetime(f\"{PREDICTION_JOB_PREFIX}_{model_name}\")\n",
+        "predict_destination_prefix = os.path.join(STAGING_BUCKET, predict_job_name)\n",
+        "\n",
+        "batch_prediction_job = model.batch_predict(\n",
+        "    job_display_name=predict_job_name,\n",
+        "    gcs_source=test_jsonl_path,\n",
+        "    gcs_destination_prefix=predict_destination_prefix,\n",
+        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    accelerator_count=PREDICTION_ACCELERATOR_COUNT,\n",
+        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    max_replica_count=1,\n",
+        "    service_account=batch_predict_service_account,\n",
+        ")\n",
+        "\n",
+        "batch_prediction_job.wait()\n",
+        "\n",
+        "print(batch_prediction_job.display_name)\n",
+        "print(batch_prediction_job.resource_name)\n",
+        "print(batch_prediction_job.state)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ik-XPjfx9OCE"
+      },
+      "source": [
+        "You can then read the prediction response JSONL files in the output directory:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tdkW9e5B9OU1"
+      },
+      "outputs": [],
+      "source": [
+        "# The label map file was generated from the section above (`Prepare input data for training`).\n",
+        "for file in tf.io.gfile.glob(os.path.join(predict_destination_prefix, \"*/*\")):\n",
+        "    with tf.io.gfile.GFile(file, \"r\") as f:\n",
+        "        for line in f:\n",
+        "            print_response_instance(line, label_map)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "## Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete the trained model.\n",
+        "model.delete()\n",
+        "# Delete custom and hpt jobs.\n",
+        "if data_converter_custom_job.list(filter=f'display_name=\"{data_converter_job_name}\"'):\n",
+        "    data_converter_custom_job.delete()\n",
+        "if train_hpt_job.list(filter=f'display_name=\"{train_job_name}\"'):\n",
+        "    train_hpt_job.delete()\n",
+        "if model_export_custom_job.list(filter=f'display_name=\"{model_export_job_name}\"'):\n",
+        "    model_export_custom_job.delete()\n",
+        "if batch_prediction_job.list(filter=f'display_name=\"{predict_job_name}\"'):\n",
+        "    batch_prediction_job.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_movinet_clip_classification.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pipeline_templates_bert.ipynb

@@ -321,10 +321,8 @@
       },
       "outputs": [],
       "source": [
-        "IS_COLAB = False\n",
         "# from google.colab import auth\n",
-        "# auth.authenticate_user()\n",
-        "# IS_COLAB=True"
+        "# auth.authenticate_user()"
       ]
     },
     {
@@ -409,6 +407,9 @@
       },
       "outputs": [],
       "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
         "if (\n",
         "    SERVICE_ACCOUNT == \"\"\n",
         "    or SERVICE_ACCOUNT is None\n",
@@ -471,6 +472,7 @@
       "outputs": [],
       "source": [
         "import json\n",
+        "import os\n",
         "\n",
         "import google.cloud.aiplatform as aiplatform\n",
         "from kfp.registry import RegistryClient"
@@ -498,6 +500,33 @@
         "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2d242773d707"
+      },
+      "source": [
+        "### Enable Artifact Registry API\n",
+        "You must enable the Artifact Registry API service for your project.\n",
+        "\n",
+        "<a href=\"https://cloud.google.com/artifact-registry/docs/enable-service\">Learn more about Enabling service</a>."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "162b5e8883c2"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud services enable artifactregistry.googleapis.com\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -506,9 +535,7 @@
       "source": [
         "## Create repo in Artifact Registry\n",
         "\n",
-        "First, you create your own (user-defined) repository in the `Artifact Registry`. You use this repository to upload and retrieve your pipeline templates.\n",
-        "\n",
-        "The name of your repo is `quickstart-kfp-repo`"
+        "First, you create your own (user-defined) repository in the `Artifact Registry`. You use this repository to upload and retrieve your pipeline templates."
       ]
     },
     {
@@ -519,7 +546,7 @@
       },
       "outputs": [],
       "source": [
-        "REPO_NAME = \"quickstart-kfp-repo\"\n",
+        "REPO_NAME = \"my-docker-repo-unique\"\n",
         "\n",
         "! gcloud artifacts repositories create {REPO_NAME} --location={REGION} --repository-format=KFP"
       ]
@@ -833,7 +860,9 @@
         "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
         "    ! gsutil rm -r $BUCKET_URI\n",
         "\n",
-        "! rm -rf custom custom.tar.gz"
+        "! rm -rf custom custom.tar.gz\n",
+        "\n",
+        "! gcloud artifacts repositories delete $REPO_NAME --project {PROJECT_ID} --location {REGION} --quiet"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pipeline_templates_t5x.ipynb

@@ -319,10 +319,8 @@
       },
       "outputs": [],
       "source": [
-        "IS_COLAB = False\n",
         "# from google.colab import auth\n",
-        "# auth.authenticate_user()\n",
-        "# IS_COLAB=True"
+        "# auth.authenticate_user()"
       ]
     },
     {
@@ -407,6 +405,9 @@
       },
       "outputs": [],
       "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
         "if (\n",
         "    SERVICE_ACCOUNT == \"\"\n",
         "    or SERVICE_ACCOUNT is None\n",
@@ -469,6 +470,7 @@
       "outputs": [],
       "source": [
         "import json\n",
+        "import os\n",
         "\n",
         "import google.cloud.aiplatform as aiplatform\n",
         "from kfp.registry import RegistryClient"
@@ -496,6 +498,33 @@
         "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2d242773d707"
+      },
+      "source": [
+        "### Enable Artifact Registry API\n",
+        "You must enable the Artifact Registry API service for your project.\n",
+        "\n",
+        "<a href=\"https://cloud.google.com/artifact-registry/docs/enable-service\">Learn more about Enabling service</a>."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "162b5e8883c2"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud services enable artifactregistry.googleapis.com\n",
+        "\n",
+        "if os.getenv(\"IS_TESTING\"):\n",
+        "    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector\n",
+        "    ! gcloud components update --quiet"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -504,9 +533,7 @@
       "source": [
         "## Create repo in Artifact Registry\n",
         "\n",
-        "First, you create your own (user-defined) repository in the `Artifact Registry`. You use this repository to upload and retreive your pipeline templates.\n",
-        "\n",
-        "The name of your repo is `quickstart-kfp-repo`"
+        "First, you create your own (user-defined) repository in the `Artifact Registry`. You use this repository to upload and retreive your pipeline templates."
       ]
     },
     {
@@ -517,7 +544,7 @@
       },
       "outputs": [],
       "source": [
-        "REPO_NAME = \"quickstart-kfp-repo\"\n",
+        "REPO_NAME = \"my-docker-repo-unique\"\n",
         "\n",
         "! gcloud artifacts repositories create {REPO_NAME} --location={REGION} --repository-format=KFP"
       ]
@@ -822,7 +849,7 @@
       },
       "outputs": [],
       "source": [
-        "delete_bucket = True\n",
+        "delete_bucket = False\n",
         "\n",
         "endpoint.undeploy_all()\n",
         "endpoint.delete()\n",
@@ -831,7 +858,9 @@
         "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
         "    ! gsutil rm -r $BUCKET_URI\n",
         "\n",
-        "! rm -rf custom custom.tar.gz"
+        "! rm -rf custom custom.tar.gz\n",
+        "\n",
+        "! gcloud artifacts repositories delete $REPO_NAME --project {PROJECT_ID} --location {REGION} --quiet"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_proprietary_image_classification.ipynb

@@ -24,7 +24,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TirJ-SGQseby"
@@ -55,7 +54,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "dwGLvtIeECLK"
@@ -67,7 +65,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -107,7 +104,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "KEukV6uRk_S3"
@@ -138,7 +134,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "BF1j6f9HApxa"
@@ -196,7 +191,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "n6IFz75WGCam"
@@ -229,16 +223,10 @@
         "# Evaluation constants.\n",
         "EVALUATION_METRIC = \"accuracy\"\n",
         "\n",
-        "# Prediction constants.\n",
-        "# The example in this notebook uses optimized tensorflow runtime dockers.\n",
-        "# You can adjust accelerator types and machine types to get faster predictions.\n",
-        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
-        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
         "DEPLOY_JOB_PREFIX = \"deploy\""
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ZZFPe_GezXg8"
@@ -301,7 +289,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "Q149N3V6Uynm"
@@ -330,7 +317,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "8yfBZ1_8VZvq"
@@ -363,7 +349,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "RB_xY9ipr7ZU"
@@ -389,7 +374,7 @@
         "- `model_type`: The type of model for deployment.\n",
         "  - `EFFICIENTNET`: A model that is available in Vertex Model Garden image classification training with customizable hyperparameters. Best tailored to be used within Google Cloud, and cannot be exported externally.\n",
         "  - `MAXVIT`: A model that is available in Vertex Model Garden image classification training with customizable hyperparameters. Best tailored to be used within Google Cloud, and cannot be exported externally.\n",
-        "  - `COCA`: A model that is available in Vertex Model Garden image classification training with customizable hyperparameters. Best tailored to be used within Google Cloud, and cannot be exported externally.\n",
+        "  - `VIT`: A model that is available in Vertex Model Garden image classification training with customizable hyperparameters. Best tailored to be used within Google Cloud, and cannot be exported externally.\n",
         "- `checkpoint_name`: Optional. The field is reserved for Model Garden model training, based on the provided pre-trained model checkpoint.\n",
         "- `trainer_config`: Optional. The field is usually used together with the Model Garden model training when passing the customized configs for the trainer.\n",
         "\n",
@@ -404,6 +389,9 @@
         "    'accelerator_count': '1',\n",
         "  }\n",
         "```\n",
+        "  The global_batch_size should be divisible by accelerator_count.\n",
+        "  Supported values for optimizer_type are 'sgd', 'adam', 'adamw', 'lamb', 'rmsprop', 'lars', 'adagrad', and 'slide'.\n",
+        "  Supported values for accelerator_count are '1', '2', '4', and '8'.\n",
         "- `metric_spec`: Dictionary representing metrics to optimize. The dictionary key is the `metric_id`, which is reported by your training job, with possible values being ('loss', 'accuracy') and the dictionary value is the optimization goal of the metric ('minimize' or 'maximize').\n",
         "For example:  `metric_spec = {'loss': 'minimize', 'accuracy': 'maximize'}`\n",
         "- `parameter_spec`: Dictionary representing parameters to optimize. The dictionary key is the `metric_id`, which is passed into your training job as a command line key word argument, and the dictionary value is the parameter\n",
@@ -457,7 +445,7 @@
         "METRIC_SPEC_VALUE = \"maximize\"\n",
         "SEARCH_ALGORITHM = \"random\"\n",
         "MEASUREMENT_SELECTION = \"best\"\n",
-        "MODEL_TYPE = \"COCA\"  # @param {type:\"string\"} one of the values [\"COCA\", \"MAXVIT\", \"EFFICIENTNET\"]\n",
+        "MODEL_TYPE = \"MAXVIT\"  # @param {type:\"string\"} one of the values [\"MAXVIT\", \"EFFICIENTNET\", \"VIT\"]\n",
         "\n",
         "job = aiplatform.AutoMLImageTrainingJob(\n",
         "    display_name=get_job_name_with_datetime(TRAINING_JOB_PREFIX),\n",
@@ -478,7 +466,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "HwcCjwlBTQIz"
@@ -523,7 +510,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "g0BGaofgsMsy"
@@ -544,16 +530,15 @@
       "outputs": [],
       "source": [
         "# @title Deploy model from Model Registry\n",
+        "# Model does not support dedicated deployment resources.\n",
+        "# An n1-standard-4 machine with 1 P100 GPU will be used.\n",
         "\n",
         "deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
         "print(\"The deployed job name is: \", deploy_model_name)\n",
         "\n",
         "endpoint = model.deploy(\n",
         "    deployed_model_display_name=deploy_model_name,\n",
-        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
         "    traffic_split={\"0\": 100},\n",
-        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
-        "    accelerator_count=1,\n",
         "    min_replica_count=1,\n",
         "    max_replica_count=1,\n",
         ")\n",
@@ -572,7 +557,7 @@
       "source": [
         "# @title Run online predictions\n",
         "\n",
-        "# test image file path from a GCS bucket\n",
+        "# test image file path from a Cloud Storage bucket\n",
         "test_filepath = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "with tf.io.gfile.GFile(test_filepath, \"rb\") as f:\n",
@@ -589,7 +574,66 @@
       ]
     },
     {
-      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EYgGK2zwtAdO"
+      },
+      "source": [
+        "# Run batch predictions\n",
+        "Now that your Model resource is trained, you can make a batch prediction by invoking the `batch_predict()` method, with the following parameters:\n",
+        "\n",
+        "* `job_display_name`: The human readable name for the batch prediction job.\n",
+        "* `gcs_source`: A jsonl file path from a Cloud Storage bucket, with a list of one or more images.\n",
+        "* `gcs_destination_prefix`: The Cloud Storage location for storing the batch prediction resuls.\n",
+        "* `sync`: If set to True, the call block while waiting for the asynchronous batch job to complete."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1z9d8oPXtD95"
+      },
+      "outputs": [],
+      "source": [
+        "# A jsonl file path from a Cloud Storage bucket, with all the to-be-predicted images.\n",
+        "gcs_source = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "batch_predict_job = model.batch_predict(\n",
+        "    job_display_name=get_job_name_with_datetime(\"flowers_bp\"),\n",
+        "    gcs_source=gcs_source,\n",
+        "    gcs_destination_prefix=f\"gs://{BUCKET_URI}\",\n",
+        "    sync=False,\n",
+        ")\n",
+        "print(batch_predict_job)\n",
+        "\n",
+        "# Wait for the batch prediction job to finish\n",
+        "batch_predict_job.wait()\n",
+        "\n",
+        "\n",
+        "# Get the batch prediction results\n",
+        "import json\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "\n",
+        "bp_iter_outputs = batch_predict_job.iter_outputs()\n",
+        "\n",
+        "prediction_results = list()\n",
+        "for blob in bp_iter_outputs:\n",
+        "    if blob.name.split(\"/\")[-1].startswith(\"prediction\"):\n",
+        "        prediction_results.append(blob.name)\n",
+        "\n",
+        "tags = list()\n",
+        "for prediction_result in prediction_results:\n",
+        "    gfile_name = f\"gs://{bp_iter_outputs.bucket.name}/{prediction_result}\"\n",
+        "    with tf.io.gfile.GFile(name=gfile_name, mode=\"r\") as gfile:\n",
+        "        for line in gfile.readlines():\n",
+        "            line = json.loads(line)\n",
+        "            print(line)\n",
+        "            break"
+      ]
+    },
+    {
       "cell_type": "markdown",
       "metadata": {
         "id": "kkH2nrpdp4sp"
@@ -606,12 +650,22 @@
       },
       "outputs": [],
       "source": [
+        "# Delete the dataset.\n",
+        "if \"dataset\" in globals():\n",
+        "    dataset.delete()\n",
+        "\n",
         "# Undeploy model and delete endpoint.\n",
-        "endpoint.undeploy_all()\n",
-        "endpoint.delete(force=True)\n",
+        "if \"endpoint\" in globals():\n",
+        "    endpoint.undeploy_all()\n",
+        "    endpoint.delete(force=True)\n",
         "\n",
         "# Delete models.\n",
-        "model.delete()"
+        "if \"model\" in globals():\n",
+        "    model.delete()\n",
+        "\n",
+        "# Delete the batch predictio job.\n",
+        "if \"batch_prediction_job\" in globals():\n",
+        "    batch_predict_job.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_proprietary_image_object_detection.ipynb

@@ -24,7 +24,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TirJ-SGQseby"
@@ -55,7 +54,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "dwGLvtIeECLK"
@@ -67,7 +65,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -107,7 +104,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "KEukV6uRk_S3"
@@ -138,7 +134,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "BF1j6f9HApxa"
@@ -181,11 +176,11 @@
         "# You can choose a region from https://cloud.google.com/about/locations.\n",
         "# Only regions prefixed by \"us\", \"europe\", or \"asia\" are supported.\n",
         "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
-        "REGION_PREFIX = REGION.split('-')[0]\n",
+        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
         "assert REGION_PREFIX in (\n",
         "    \"us\",\n",
         "    \"europe\",\n",
-        "    \"asia\"\n",
+        "    \"asia\",\n",
         "), f'{REGION} is not supported. It must be prefixed by \"us\", \"europe\", or \"asia\".'\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
@@ -196,7 +191,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "n6IFz75WGCam"
@@ -226,15 +220,10 @@
         "# Evaluation constants.\n",
         "EVALUATION_METRIC = \"AP50\"\n",
         "\n",
-        "# Prediction constants.\n",
-        "# You can adjust accelerator types and machine types to get faster predictions.\n",
-        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
-        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
         "DEPLOY_JOB_PREFIX = \"deploy\""
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ZZFPe_GezXg8"
@@ -357,7 +346,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "nZLVI9TtUuif"
@@ -386,7 +374,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "SZEdBfNZUxQn"
@@ -419,7 +406,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "RB_xY9ipr7ZU"
@@ -455,10 +441,13 @@
         "    'optimizer_type': 'sgd',\n",
         "    'optimizer_momentum': '0.9',\n",
         "    'train_steps': '10000',\n",
-        "    'accelerator_count': '1',\n",
+        "    'accelerator_count': '2',\n",
         "    'anchor_size': '8',\n",
         "  }\n",
         "```\n",
+        "  The global_batch_size should be divisible by accelerator_count.\n",
+        "  Supported values for optimizer_type are 'sgd', 'adam', 'adamw', 'lamb', 'rmsprop', 'lars', 'adagrad', and 'slide'.\n",
+        "  Supported values for accelerator_count are '2', '4', and '8'.\n",
         "- `metric_spec`: Dictionary representing metrics to optimize. The dictionary key is the metric_id, which is reported by your training job, with possible values being ('loss', 'AP50') and the dictionary value is the optimization goal of the metric('minimize' or 'maximize').\n",
         "For example:  `metric_spec = {'loss': 'minimize', 'AP50': 'maximize'}`\n",
         "- `parameter_spec`:Dictionary representing parameters to optimize. The dictionary key is the `metric_id`, which is passed into your training job as a command line key word argument, and the dictionary value is the parameter\n",
@@ -506,7 +495,7 @@
         "    \"global_batch_size\": \"8\",\n",
         "    \"learning_rate\": \"0.001\",\n",
         "    \"train_steps\": \"10000\",\n",
-        "    \"accelerator_count\": \"1\",\n",
+        "    \"accelerator_count\": \"2\",\n",
         "}\n",
         "METRIC_SPEC_KEY = \"AP50\"\n",
         "METRIC_SPEC_VALUE = \"maximize\"\n",
@@ -550,7 +539,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "HwcCjwlBTQIz"
@@ -595,7 +583,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "g0BGaofgsMsy"
@@ -616,16 +603,15 @@
       "outputs": [],
       "source": [
         "# @title Deploy model from Model Registry\n",
+        "# Model does not support dedicated deployment resources.\n",
+        "# An n1-standard-4 machine with 1 P100 GPU will be used.\n",
         "\n",
         "deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
         "print(\"The deployed job name is: \", deploy_model_name)\n",
         "\n",
         "endpoint = model.deploy(\n",
         "    deployed_model_display_name=deploy_model_name,\n",
-        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
         "    traffic_split={\"0\": 100},\n",
-        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
-        "    accelerator_count=1,\n",
         "    min_replica_count=1,\n",
         "    max_replica_count=1,\n",
         ")\n",
@@ -644,7 +630,7 @@
       "source": [
         "# @title Run online predictions\n",
         "\n",
-        "# test image file path from a GCS bucket\n",
+        "# test image file path from a Cloud Storage bucket\n",
         "test_filepath = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "with tf.io.gfile.GFile(test_filepath, \"rb\") as f:\n",
@@ -661,7 +647,66 @@
       ]
     },
     {
-      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "MNPNJyUnrKY3"
+      },
+      "source": [
+        "# Run batch predictions\n",
+        "Now that your Model resource is trained, you can make a batch prediction by invoking the `batch_predict()` method, with the following parameters:\n",
+        "\n",
+        "* `job_display_name`: The human readable name for the batch prediction job.\n",
+        "* `gcs_source`: A jsonl file path from a Cloud Storage bucket, with a list of one or more images.\n",
+        "* `gcs_destination_prefix`: The Cloud Storage location for storing the batch prediction resuls.\n",
+        "* `sync`: If set to True, the call block while waiting for the asynchronous batch job to complete."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e-MnutG5rNKR"
+      },
+      "outputs": [],
+      "source": [
+        "# A jsonl file path from a Cloud Storage bucket, with all the to-be-predicted images.\n",
+        "gcs_source = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "batch_predict_job = model.batch_predict(\n",
+        "    job_display_name=get_job_name_with_datetime(\"flowers_bp\"),\n",
+        "    gcs_source=gcs_source,\n",
+        "    gcs_destination_prefix=f\"gs://{BUCKET_URI}\",\n",
+        "    sync=False,\n",
+        ")\n",
+        "print(batch_predict_job)\n",
+        "\n",
+        "# Wait for the batch prediction job to finish\n",
+        "batch_predict_job.wait()\n",
+        "\n",
+        "\n",
+        "# Get the batch prediction results\n",
+        "import json\n",
+        "\n",
+        "import tensorflow as tf\n",
+        "\n",
+        "bp_iter_outputs = batch_predict_job.iter_outputs()\n",
+        "\n",
+        "prediction_results = list()\n",
+        "for blob in bp_iter_outputs:\n",
+        "    if blob.name.split(\"/\")[-1].startswith(\"prediction\"):\n",
+        "        prediction_results.append(blob.name)\n",
+        "\n",
+        "tags = list()\n",
+        "for prediction_result in prediction_results:\n",
+        "    gfile_name = f\"gs://{bp_iter_outputs.bucket.name}/{prediction_result}\"\n",
+        "    with tf.io.gfile.GFile(name=gfile_name, mode=\"r\") as gfile:\n",
+        "        for line in gfile.readlines():\n",
+        "            line = json.loads(line)\n",
+        "            print(line)\n",
+        "            break"
+      ]
+    },
+    {
       "cell_type": "markdown",
       "metadata": {
         "id": "frcGP5HFX1XN"
@@ -678,12 +723,22 @@
       },
       "outputs": [],
       "source": [
+        "# Delete the dataset.\n",
+        "if \"dataset\" in globals():\n",
+        "    dataset.delete()\n",
+        "\n",
         "# Undeploy model and delete endpoint.\n",
-        "endpoint.undeploy_all()\n",
-        "endpoint.delete(force=True)\n",
+        "if \"endpoint\" in globals():\n",
+        "    endpoint.undeploy_all()\n",
+        "    endpoint.delete(force=True)\n",
         "\n",
         "# Delete models.\n",
-        "model.delete()"
+        "if \"model\" in globals():\n",
+        "    model.delete()\n",
+        "\n",
+        "# Delete the batch predictio job.\n",
+        "if \"batch_prediction_job\" in globals():\n",
+        "    batch_predict_job.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_bart_large_cnn.ipynb

@@ -208,7 +208,7 @@
       "source": [
         "# The pre-built serving docker image.\n",
         "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_blip2.ipynb

@@ -205,7 +205,7 @@
       "source": [
         "# The pre-built serving docker image.\n",
         "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_blip_vqa.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_clip.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb

@@ -0,0 +1,469 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Code LLaMA\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a> (A Python-3 CPU notebook is recommended)\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying pretrained Code LLaMA models in Vertex AI.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy pre-trained Code LLaMA models with [vLLM](https://github.com/vllm-project/vllm) with best serving throughput\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/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": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ioensNKM8ned"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands for Colab and skip this section if you are using Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "    ! pip3 install ipython pandas[output_formatting] google-cloud-language==2.10.0\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "    # Install gdown for downloading example training images.\n",
+        "    ! pip3 install gdown\n",
+        "\n",
+        "    # Restart the notebook kernel after installs.\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bb7adab99e41"
+      },
+      "source": [
+        "### Setup Google Cloud project\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API, Compute Engine API and Cloud Natural Language API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component,language.googleapis.com).\n",
+        "\n",
+        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c460088b873"
+      },
+      "source": [
+        "Fill following variables for experiments environment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# Cloud project id.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The region you want to launch jobs in.\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The Cloud Storage bucket for storing experiments output.\n",
+        "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
+        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud services enable language.googleapis.com\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(GCS_BUCKET, \"staging\")\n",
+        "MODEL_BUCKET = os.path.join(GCS_BUCKET, \"code-llama\")\n",
+        "\n",
+        "# The service account looks like:\n",
+        "# '@.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "# The service account for deploying fine tuned model.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "from google.colab import auth\n",
+        "\n",
+        "auth.authenticate_user(project_id=PROJECT_ID)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e828eb320337"
+      },
+      "source": [
+        "### Initialize Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "12cd25839741"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built serving docker images.\n",
+        "VLLM_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "354da31189dc"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str):\n",
+        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
+        "    jobs in Vertex AI.\n",
+        "    \"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name,\n",
+        "    model_id,\n",
+        "    service_account,\n",
+        "    machine_type=\"n1-standard-8\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count=1,\n",
+        "):\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    vllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        \"--gpu-memory-utilization=0.95\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ivs2RK093c8X"
+      },
+      "source": [
+        "## Access pretrained Code LLaMA models\n",
+        "The original models from Meta are converted into the HuggingFace format for serving in Vertex AI.\n",
+        "After clicking the agreement of Code LLaMA model card in Vertex AI Model Garden, a Cloud Storage bucket will be shared to access pretrained Code LLaMA models."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jwn4PcTf4EMt"
+      },
+      "outputs": [],
+      "source": [
+        "VERTEX_AI_MODEL_GARDEN_CODE_LLAMA = (\n",
+        "    \"\"  # This path will be shared once click the agreement in Vertex AI Model Garden.\n",
+        ")\n",
+        "assert (\n",
+        "    VERTEX_AI_MODEL_GARDEN_CODE_LLAMA\n",
+        "), \"Please click the agreement of Code LLaMA in Vertex AI Model Garden, and get the GCS path of Code LLaMA model artifacts.\"\n",
+        "print(\n",
+        "    \"Copy Code LLaMA model artifacts from\",\n",
+        "    VERTEX_AI_MODEL_GARDEN_CODE_LLAMA,\n",
+        "    \"to \",\n",
+        "    MODEL_BUCKET,\n",
+        ")\n",
+        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_CODE_LLAMA/* $MODEL_BUCKET"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "model_name = \"CodeLlama-7b-Instruct-hf\"  # @param [\"CodeLlama-7b-hf\", \"CodeLlama-7b-Python-hf\", \"CodeLlama-7b-Instruct-hf\", \"CodeLlama-13b-hf\", \"CodeLlama-13b-Python-hf\", \"CodeLlama-13b-Instruct-hf\", \"CodeLlama-34b-hf\", \"CodeLlama-34b-Python-hf\", \"CodeLlama-34b-Instruct-hf\"]\n",
+        "model_id = os.path.join(MODEL_BUCKET, model_name)\n",
+        "print(model_id)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy Pretrained Code LLaMA with vLLM\n",
+        "\n",
+        "This section deploys prebuilt OpenLLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The vLLM project is an highly optimized LLM serving framework which can increase serving throughput a lot. The higher QPS you have, the more benefits you get using vLLM."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "03d504bcd60b"
+      },
+      "outputs": [],
+      "source": [
+        "# Since all models are stored in bfloat16 precision, A100 is needed for serving.\n",
+        "\n",
+        "# Sets A100 (40G) to deploy 7B.\n",
+        "machine_type = \"a2-highgpu-1g\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Sets A100 (40G) to deploy 13B.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 1\n",
+        "\n",
+        "# Sets 2 A100 (40G) to deploy 34B models.\n",
+        "# machine_type = \"a2-highgpu-2g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 2\n",
+        "\n",
+        "model_vllm, endpoint_vllm = deploy_model_vllm(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"code-llama-serve-vllm\"),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sGKIjgmDFRW2"
+      },
+      "source": [
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the $MODEL_BUCKET after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3f5a1e1de60d"
+      },
+      "outputs": [],
+      "source": [
+        "instance = {\n",
+        "    \"prompt\": \"import argparse\",\n",
+        "    \"n\": 1,\n",
+        "    \"max_tokens\": 200,\n",
+        "}\n",
+        "response = endpoint_vllm.predict(instances=[instance])\n",
+        "print(response.predictions[0])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint_vllm.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_vllm.delete()\n",
+        "\n",
+        "# Uncomment below to delete all artifacts\n",
+        "# !gsutil -m rm -r $STAGING_BUCKET $MODEL_BUCKET"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_codellama.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_controlnet.ipynb

@@ -259,10 +259,10 @@
       "outputs": [],
       "source": [
         "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb

@@ -44,7 +44,7 @@
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>                                                                                               <td>\n",
+        "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",

notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb

@@ -0,0 +1,902 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Falcon Instruct (PEFT)\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a> (A Python-3 GPU notebook is recommended)\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates running inferences locally with prebuilt Falcon Instruct models, deploying prebuilt Falcon Instruct models, finetuning and deploying Falcon Instruct models with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)), and evaluating PEFT-finetuned Falcon Instruct models in Vertex AI.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Run inferences locally on prebuilt Falcon Instruct models\n",
+        "- Deploy prebuilt Falcon Instruct models\n",
+        "- Finetune and deploy Falcon Instruct models with PEFT\n",
+        "- Evaluate PEFT-finetuned Falcon Instruct models\n",
+        "\n",
+        "| Models | LoRA |\n",
+        "| :- | :- |\n",
+        "| [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct) | Y |\n",
+        "| [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) | Y |\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/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": "264c07757582"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
+        "\n",
+        "Running inferences locally with Falcon Instruct models requires a GPU."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ioensNKM8ned"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands for Colab and skip this section if you are using Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "    # Install gdown for downloading example training images.\n",
+        "    ! pip3 install gdown\n",
+        "\n",
+        "    # Restart the notebook kernel after installs.\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "! pip3 install transformers==4.31.0\n",
+        "! pip3 install einops==0.6.1\n",
+        "! pip3 install accelerate==0.21.0"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bb7adab99e41"
+      },
+      "source": [
+        "### Setup Google Cloud project\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c460088b873"
+      },
+      "source": [
+        "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# Cloud project id.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The region you want to launch jobs in.\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The Cloud Storage bucket for storing experiments output with gs:// prefix.\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
+        "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
+        "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
+        "\n",
+        "# The service account looks like:\n",
+        "# '@.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "# The service account for deploying fine tuned model.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e828eb320337"
+      },
+      "source": [
+        "### Initialize Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "12cd25839741"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built training, serving and evaluation docker images.\n",
+        "TRAIN_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        ")\n",
+        "PREDICTION_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        ")\n",
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "354da31189dc"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
+        "    jobs in Vertex AI.\n",
+        "    \"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def deploy_model(\n",
+        "    model_name: str,\n",
+        "    base_model_id: str,\n",
+        "    finetuned_lora_model_path: str,\n",
+        "    service_account: str,\n",
+        "    task: str,\n",
+        "    machine_type: str = \"n1-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count: int = 1,\n",
+        ") -> tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    serving_env = {\n",
+        "        \"BASE_MODEL_ID\": base_model_id,\n",
+        "        \"TASK\": task,\n",
+        "    }\n",
+        "    if finetuned_lora_model_path:\n",
+        "        serving_env[\"FINETUNED_LORA_MODEL_PATH\"] = finetuned_lora_model_path\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predictions/peft_serving\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "06e73cb3f412"
+      },
+      "source": [
+        "## Run inferences locally with prebuilt Falcon Instruct models\n",
+        "\n",
+        "You will need at least 16GB of memory to swiftly run inference with Falcon-7B-Instruct."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3ea64305957f"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "import transformers\n",
+        "from transformers import AutoTokenizer\n",
+        "\n",
+        "model = \"tiiuae/falcon-7b-instruct\"\n",
+        "\n",
+        "tokenizer = AutoTokenizer.from_pretrained(model)\n",
+        "pipeline = transformers.pipeline(\n",
+        "    \"text-generation\",\n",
+        "    model=model,\n",
+        "    tokenizer=tokenizer,\n",
+        "    torch_dtype=torch.bfloat16,\n",
+        "    trust_remote_code=True,\n",
+        "    device_map=\"auto\",\n",
+        ")\n",
+        "sequences = pipeline(\n",
+        "    \"Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Girafatron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\\nDaniel: Hello, Girafatron!\\nGirafatron:\",\n",
+        "    max_length=200,\n",
+        "    do_sample=True,\n",
+        "    top_k=10,\n",
+        "    num_return_sequences=1,\n",
+        "    eos_token_id=tokenizer.eos_token_id,\n",
+        ")\n",
+        "for seq in sequences:\n",
+        "    print(f\"Result: {seq['generated_text']}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy prebuilt Falcon Instruct models\n",
+        "\n",
+        "This section deploys prebuilt Falcon Instruct models on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The peak GPU memory usages for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) are ~15.5G and ~84G separately with the default settings. Please adjust the machine type, accelerator type and accelerator count accordingly. We use V100 in deployments as an example. Please use A100 (40G) or A100 (80G) to get better inferences."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2MjaORIIFDVu"
+      },
+      "source": [
+        "Set the prebuilt model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "prebuilt_model_id = \"tiiuae/falcon-7b-instruct\"  # @param [\"tiiuae/falcon-7b-instruct\", \"tiiuae/falcon-40b-instruct\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dHFW7yvjaVFV"
+      },
+      "source": [
+        "We use the PEFT serving images to deploy prebuilt Falcon Instruct models, by setting finetuning LoRA model paths as empty."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Uak1pyEeExYM"
+      },
+      "outputs": [],
+      "source": [
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/predictions/configure-compute\n",
+        "\n",
+        "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "# If A100 is not available, you may deploy tiiuae/falcon-40b-instruct with\n",
+        "#  multiple V100s. Please keep in mind that the efficiency of serving with\n",
+        "#  multiple V100s is inferior to that of serving with A100s.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 8  # for tiiuae/falcon-40b-instruct\n",
+        "\n",
+        "# Sets A100 (40G) to deploy tiiuae/falcon-7b-instruct or tiiuae/falcon-40b-instruct.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 4  # for tiiuae/falcon-40b-instruct\n",
+        "\n",
+        "# Sets A100 (80G) to deploy falcon-40b-instruct models for faster inferences.\n",
+        "# machine_type = \"a2-ultragpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 2\n",
+        "\n",
+        "model_without_peft, endpoint_without_peft = deploy_model(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"falcon-instruct-serve\"),\n",
+        "    base_model_id=prebuilt_model_id,\n",
+        "    finetuned_lora_model_path=\"\",  # This will avoid override finetuning models.\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    task=\"instruct-lora\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")\n",
+        "print(\"endpoint_name:\", endpoint_without_peft.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sGKIjgmDFRW2"
+      },
+      "source": [
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "Example:\n",
+        "\n",
+        "```\n",
+        "Human: What is a car?\n",
+        "Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rDHsCOqvFYBi"
+      },
+      "outputs": [],
+      "source": [
+        "# # Loads an existing endpoint as below.\n",
+        "# endpoint_name = endpoint_without_peft.name\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_without_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "instances = [\n",
+        "    {\"prompt\": \"What is a car?\"},\n",
+        "]\n",
+        "response = endpoint_without_peft.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions[0]:\n",
+        "    print(prediction[\"generated_text\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e70e3519ff8b"
+      },
+      "source": [
+        "## Finetune and deploy Falcon Instruct models with PEFT\n",
+        "\n",
+        "This section demonstrates how to finetune and deploy Falcon Instruct models with PEFT LoRA."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5qCrm_kJH5cz"
+      },
+      "source": [
+        "Set the base model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "N3UBLiYrM3sU"
+      },
+      "outputs": [],
+      "source": [
+        "base_model_id = \"tiiuae/falcon-7b-instruct\"  # @param [\"tiiuae/falcon-7b-instruct\", \"tiiuae/falcon-40b-instruct\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iWGwJHqI7LMs"
+      },
+      "source": [
+        "### Finetune"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KKEYoRfiHDVv"
+      },
+      "source": [
+        "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
+        "\n",
+        "This example uses the dataset [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).\n",
+        "\n",
+        "The peak GPU memory usages are ~11G and ~34G for finetuning LoRA models for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) separately with default training parameters and the example dataset. Falcon-7b-instruct can be finetuned on 1 P100/V100, and falcon-40b-instruct can be finetuned on 1 A100 (40G)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "65467b361315"
+      },
+      "outputs": [],
+      "source": [
+        "dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
+        "\n",
+        "# Uses V100 to finetune falcon-7b-instruct.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "# Uses A100 to finetune falcon-40b-instruct.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "replica_count = 1\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Setup training job.\n",
+        "job_name = get_job_name_with_datetime(\"falcon-instruct-lora-train\")\n",
+        "train_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Pass training arguments and launch job.\n",
+        "max_steps = 10\n",
+        "train_job.run(\n",
+        "    args=[\n",
+        "        \"--task=instruct-lora\",\n",
+        "        f\"--pretrained_model_id={base_model_id}\",\n",
+        "        f\"--dataset_name={dataset_name}\",\n",
+        "        f\"--output_dir={output_dir_gcsfuse}\",\n",
+        "        \"--lora_rank=64\",\n",
+        "        \"--lora_alpha=16\",\n",
+        "        \"--lora_dropout=0.1\",\n",
+        "        \"--warmup_ratio=0.03\",\n",
+        "        f\"--max_steps={max_steps}\",\n",
+        "        \"--max_seq_length=512\",\n",
+        "        \"--learning_rate=2e-4\",\n",
+        "    ],\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    boot_disk_size_gb=500,\n",
+        ")\n",
+        "\n",
+        "print(\"Trained models were saved in: \", output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jqmCtkGnhDmp"
+      },
+      "source": [
+        "### Deploy\n",
+        "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
+        "\n",
+        "The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The peak GPU memory usages for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) with LoRA weights are ~15.5G and ~84G separately with the default settings. Please adjust the machine type, accelerator type and accelerator count accordingly. We use V100 in deployments as an example. Please use A100 (40G) or A100 (80G) to get better inferences."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf55e38815dc"
+      },
+      "outputs": [],
+      "source": [
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/predictions/configure-compute\n",
+        "\n",
+        "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "# If A100 is not available, you may deploy tiiuae/falcon-40b-instruct with\n",
+        "#  multiple V100s. Please keep in mind that the efficiency of serving with\n",
+        "#  multiple V100s is inferior to that of serving with A100s.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 8  # for tiiuae/falcon-40b-instruct\n",
+        "\n",
+        "# Sets A100 (40G) to deploy tiiuae/falcon-7b-instruct or tiiuae/falcon-40b-instruct.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 4  # for tiiuae/falcon-40b-instruct\n",
+        "\n",
+        "# Sets A100 (80G) to deploy falcon-40b-instruct models for faster inferences.\n",
+        "# machine_type = \"a2-ultragpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 2\n",
+        "\n",
+        "model_with_peft, endpoint_with_peft = deploy_model(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"falcon-instruct-peft-serve\"),\n",
+        "    base_model_id=base_model_id,\n",
+        "    finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\" + str(max_steps)),\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    task=\"instruct-lora\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")\n",
+        "print(\"endpoint_name:\", endpoint_with_peft.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "80b3fd2ace09"
+      },
+      "source": [
+        "NOTE: After the deployment succeeds, the base model weights will be downloaded one the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "Example:\n",
+        "\n",
+        "```\n",
+        "Human: What is a car?\n",
+        "Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4ab04da3ec9a"
+      },
+      "outputs": [],
+      "source": [
+        "# # Loads an existing endpoint as below.\n",
+        "# endpoint_name = endpoint_with_peft.name\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_with_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "instances = [\n",
+        "    {\"prompt\": \"What is a car?\"},\n",
+        "]\n",
+        "response = endpoint_with_peft.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions[0]:\n",
+        "    print(prediction[\"generated_text\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4ffffac5b20a"
+      },
+      "source": [
+        "## Evaluate PEFT-finetuned Falcon Instruct models\n",
+        "\n",
+        "This section demonstrates how to evaluate the Falcon Instruct models fintuned with PEFT LoRA using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob. Please reference the peak GPU memory usgaes for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
+        "\n",
+        "This example uses the dataset [HellaSwag](https://allenai.org/data/hellaswag). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "078b4d178624"
+      },
+      "outputs": [],
+      "source": [
+        "eval_dataset = \"hellaswag\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
+        "\n",
+        "# Sets V100 to deploy tiiuae/falcon-7b-instruct as an example.\n",
+        "# If A100 is not available, you may deploy tiiuae/falcon-40b-instruct with\n",
+        "#  multiple V100s. Please keep in mind that the efficiency of serving with\n",
+        "#  multiple V100s is inferior to that of serving with A100s.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 8  # for tiiuae/falcon-40b-instruct\n",
+        "\n",
+        "# Sets A100 (40G) to deploy tiiuae/falcon-7b-instruct or tiiuae/falcon-40b-instruct.\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 1  # for tiiuae/falcon-7b-instruct\n",
+        "# accelerator_count = 4  # for tiiuae/falcon-40b-instruct\n",
+        "\n",
+        "# Sets A100 (80G) to deploy falcon-40b-instruct models for faster inferences.\n",
+        "# machine_type = \"a2-ultragpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "# accelerator_count = 2\n",
+        "\n",
+        "replica_count = 1\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "job_name = get_job_name_with_datetime(prefix=\"falcon-instruct-peft-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c5e32112815e"
+      },
+      "outputs": [],
+      "source": [
+        "# Prepare evaluation script that runs the evaluation harness.\n",
+        "# We set `trust_remote_code = True` because evaluating the model requires\n",
+        "#  executing code from the model repository.\n",
+        "# We set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
+        "script_path = \"./eval_script.py\"  # @param {type:\"string\"}\n",
+        "\n",
+        "eval_command = f\"\"\"import subprocess\n",
+        "\n",
+        "\n",
+        "subprocess.call([\n",
+        "    'python',\n",
+        "    'main.py',\n",
+        "    '--model',\n",
+        "    'hf-causal-experimental',\n",
+        "    '--model_args',\n",
+        "    'pretrained={base_model_id},peft={output_dir_gcsfuse},trust_remote_code=True,use_accelerate=True,device_map_option=auto',\n",
+        "    '--tasks',\n",
+        "    '{eval_dataset}',\n",
+        "    '--output_path',\n",
+        "    '{eval_output_dir_gcsfuse}',\n",
+        "])\n",
+        "\"\"\"\n",
+        "\n",
+        "with open(script_path, \"w\") as fp:\n",
+        "    fp.write(eval_command)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "afd9305771ba"
+      },
+      "source": [
+        "### Submit evaluation CustomJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "602413896b58"
+      },
+      "outputs": [],
+      "source": [
+        "# Pass evaluation arguments and launch job.\n",
+        "eval_job = aiplatform.CustomJob.from_local_script(\n",
+        "    display_name=job_name,\n",
+        "    script_path=script_path,\n",
+        "    container_uri=EVAL_DOCKER_URI,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    base_output_dir=eval_output_dir,\n",
+        "    boot_disk_size_gb=500,\n",
+        ")\n",
+        "\n",
+        "eval_job.run()\n",
+        "\n",
+        "print(\"Evaluation results were saved in:\", eval_output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de47e182a37e"
+      },
+      "source": [
+        "### Fetch and print evaluation results"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1f15ed6d375a"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "from google.cloud import storage\n",
+        "\n",
+        "# Fetch evaluation results.\n",
+        "storage_client = storage.Client()\n",
+        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
+        "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
+        "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :]\n",
+        "blob = bucket.blob(RESULT_FILE_PATH)\n",
+        "raw_result = blob.download_as_string()\n",
+        "\n",
+        "# Print evaluation results.\n",
+        "result = json.loads(raw_result)\n",
+        "result_formatted = json.dumps(result, indent=2)\n",
+        "print(f\"Evaluation result:\\n{result_formatted}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete custom train and evaluation jobs.\n",
+        "train_job.delete()\n",
+        "eval_job.delete()\n",
+        "\n",
+        "# Undeploy models and delete endpoints.\n",
+        "endpoint_without_peft.delete(force=True)\n",
+        "endpoint_with_peft.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_without_peft.delete()\n",
+        "model_with_peft.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_falcon_instruct_peft.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_instructpix2pix.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_layoutml_document_qa.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb

@@ -0,0 +1,677 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - OpenCLIP\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates finetuning [OpenCLIP](https://github.com/mlfoundations/open_clip) with [CC3M](https://ai.google.com/research/ConceptualCaptions/download) dataset and deploying it on Vertex AI for online prediction.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Finetune the OpenCLIP model with [Vertex AI custom training](https://cloud.google.com/vertex-ai/docs/training/overview).\n",
+        "- Upload the model to [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
+        "- Deploy the model to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
+        "- Run online predictions for zero-shot image classification.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/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": "264c07757582"
+      },
+      "source": [
+        "## Setup environment\n",
+        "\n",
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d73ffa0c0b83"
+      },
+      "source": [
+        "### Colab only"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "if \"google.colab\" in str(get_ipython()):\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "\n",
+        "    # Restart the notebook kernel after installs.\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bb7adab99e41"
+      },
+      "source": [
+        "### Setup Google Cloud project\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User`,  `Storage Object Admin`, and `GCS Storage Bucket Owner roles` roles for deploying fine tuned model to Vertex AI endpoint."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c460088b873"
+      },
+      "source": [
+        "Fill the following variables for experiments environment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a12c23679315"
+      },
+      "outputs": [],
+      "source": [
+        "# Cloud project id.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The region you want to launch jobs in.\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The GCS bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
+        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The service account for deploying fine tuned model.\n",
+        "# The service account looks like:\n",
+        "# '<account_name>@<project>.iam.gserviceaccount.com'\n",
+        "# Follow step 5 above to create this account.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "12c249e14a5d"
+      },
+      "source": [
+        "### Download data to Google Cloud Storage"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fb0259bfd059"
+      },
+      "outputs": [],
+      "source": [
+        "# Install the library for downloading training data.\n",
+        "!pip install img2dataset==1.41.0"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1d6b99305e7e"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "!gcloud storage cp gs://gcc-data/Validation/GCC-1.1.0-Validation.tsv ./data.tsv  # Download list of URLs.\n",
+        "!sed -i '1s/^/caption\\turl\\n/' data.tsv  # Add column name."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f36f82f74dde"
+      },
+      "outputs": [],
+      "source": [
+        "# Download images from URLs. It takes around 2min.\n",
+        "output_folder = \"data\"\n",
+        "!img2dataset --url_list data.tsv --input_format \"tsv\"\\\n",
+        "    --output_folder {output_folder}\\\n",
+        "    --url_col \"url\" --caption_col \"caption\" --output_format webdataset\\\n",
+        "    --processes_count {os.cpu_count()} --thread_count {os.cpu_count()*4}\\\n",
+        "    --image_size 256"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "514ee0208d47"
+      },
+      "outputs": [],
+      "source": [
+        "GCS_data_dir = f\"{GCS_BUCKET}/CC3M-val-wds\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b90bbd80476a"
+      },
+      "outputs": [],
+      "source": [
+        "# Upload data to GCS.\n",
+        "!gcloud storage cp -r data gs://{GCS_data_dir}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7ffc8c1aed83"
+      },
+      "outputs": [],
+      "source": [
+        "# Get total number of samples, which is required by OpenCLIP training.\n",
+        "import json\n",
+        "\n",
+        "n_samples = 0\n",
+        "for filename in [f for f in os.listdir(output_folder) if f.endswith(\"_stats.json\")]:\n",
+        "    with open(os.path.join(output_folder, filename), \"r\") as f:\n",
+        "        n_samples += json.load(f)[\"successes\"]\n",
+        "print(n_samples)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e828eb320337"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "12cd25839741"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built training docker image. It contains training scripts and models.\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-open-clip-train\"\n",
+        "\n",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-open-clip-serve\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions for job operations and test data preparations"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "354da31189dc"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "from datetime import datetime\n",
+        "from io import BytesIO\n",
+        "\n",
+        "import matplotlib.pyplot as plt\n",
+        "import requests\n",
+        "from PIL import Image\n",
+        "\n",
+        "\n",
+        "def create_job_name(prefix):\n",
+        "    \"\"\"Create a job name string with a prefix.\"\"\"\n",
+        "    user = os.environ.get(\"USER\")\n",
+        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "    job_name = f\"{prefix}-{user}-{now}\"\n",
+        "    return job_name\n",
+        "\n",
+        "\n",
+        "def deploy_model(checkpoint, model, precision, task):\n",
+        "    \"\"\"Deploy a model to Vertex AI endpoint.\"\"\"\n",
+        "    model_name = \"openclip\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-{task}-endpoint\")\n",
+        "    serving_env = {\n",
+        "        \"CHECKPOINT\": checkpoint,\n",
+        "        \"MODEL\": model,\n",
+        "        \"PRECISION\": precision,\n",
+        "        \"TASK\": task,\n",
+        "    }\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predictions/transformers_serving\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=\"n1-standard-4\",\n",
+        "        accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "        accelerator_count=1,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def download_image(url):\n",
+        "    response = requests.get(url)\n",
+        "    return Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
+        "\n",
+        "\n",
+        "def image_to_base64(image):\n",
+        "    buffer = BytesIO()\n",
+        "    image.save(buffer, format=\"PNG\")\n",
+        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
+        "    return image_str\n",
+        "\n",
+        "\n",
+        "def plot_images(images, rows=1, cols=None):\n",
+        "    fig, axes = plt.subplots(rows, cols if cols is not None else len(images))\n",
+        "    for ax, img in zip(axes, images):\n",
+        "        ax.imshow(img)\n",
+        "        ax.axis(\"off\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e70e3519ff8b"
+      },
+      "source": [
+        "## Fine tune the model\n",
+        "\n",
+        "This section fine tunes the OpenCLIP model with [CC3M](https://ai.google.com/research/ConceptualCaptions/download) dataset validation split for the image-text pre-training. It loads the pretrained checkpoint by OpenAI. You use a small model of ***RN50*** here; check [this list](https://github.com/mlfoundations/open_clip#pretrained-model-interface) for other options supported by OpenAI checkpoints, such as ***RN50, ViT-B-32, etc***.\n",
+        "\n",
+        "One `n1-standard-4` machine with 1 `NVIDIA_TESLA_V100` is required to run the fine-tuning job. The fine-tuning job takes about 3min to complete training for 2 epochs.\n",
+        "\n",
+        "The fine-tuned model will be saved after the job finishs and it can then be loaded for inference."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "55dabb1b02e3"
+      },
+      "outputs": [],
+      "source": [
+        "machine_type = \"n1-standard-4\"\n",
+        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
+        "num_gpus = 1\n",
+        "\n",
+        "job_name = create_job_name(\"openclip\")\n",
+        "\n",
+        "model_name = \"RN50\"\n",
+        "precision = \"amp\"\n",
+        "\n",
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name, container_uri=TRAIN_DOCKER_URI, command=[\"torchrun\"]\n",
+        ")\n",
+        "\n",
+        "job.run(\n",
+        "    args=[\n",
+        "        f\"--nproc_per_node={num_gpus}\",\n",
+        "        \"-m\",\n",
+        "        \"training.main\",\n",
+        "        f\"--name={job_name}\",\n",
+        "        f\"--logs=/gcs/{GCS_BUCKET}\",\n",
+        "        f\"--train-data=/gcs/{GCS_data_dir}/{{00000..00001}}.tar\",\n",
+        "        f\"--train-num-samples={n_samples}\",\n",
+        "        \"--dataset-type=webdataset\",\n",
+        "        \"--batch-size=32\",\n",
+        "        \"--precision=amp\",\n",
+        "        \"--workers=8\",\n",
+        "        \"--dataset-resampled\",\n",
+        "        \"--save-frequency=2\",\n",
+        "        \"--epochs=2\",\n",
+        "        f\"--model={model_name}\",\n",
+        "        \"--pretrained=openai\",\n",
+        "        \"--save-most-recent\",\n",
+        "    ],\n",
+        "    boot_disk_size_gb=600,\n",
+        "    replica_count=1,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=gpu_type,\n",
+        "    accelerator_count=num_gpus,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bf7f82732e61"
+      },
+      "source": [
+        "## Upload and Deploy models\n",
+        "\n",
+        "This section uploads the fine-tuned model to Model Registry and deploys it on the Endpoint.\n",
+        "\n",
+        "One `n1-standard-4` machine with 1 `NVIDIA_TESLA_V100` is required to deploy OpenCLIP model.\n",
+        "\n",
+        "The model deployment step will take ~20 minutes to complete."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6a2cf6e84b10"
+      },
+      "outputs": [],
+      "source": [
+        "# Prepare image samples\n",
+        "img_diagram = download_image(\n",
+        "    \"https://raw.githubusercontent.com/mlfoundations/open_clip/main/docs/CLIP.png\"\n",
+        ")\n",
+        "img_cat = download_image(\n",
+        "    \"https://upload.wikimedia.org/wikipedia/commons/thumb/9/9d/Calicocats2222.jpg/220px-Calicocats2222.jpg\"\n",
+        ")\n",
+        "plot_images([img_diagram, img_cat])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5d63ffdd5f09"
+      },
+      "source": [
+        "#### Zero-shot image classification"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf55e38815dc"
+      },
+      "outputs": [],
+      "source": [
+        "model, endpoint = deploy_model(\n",
+        "    checkpoint=f\"gs://{GCS_BUCKET}/{job_name}/checkpoints/epoch_latest.pt\",\n",
+        "    model=model_name,\n",
+        "    precision=precision,\n",
+        "    task=\"zero-shot-image-classification\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "80b3fd2ace09"
+      },
+      "source": [
+        "NOTE: The model weights are downloaded after the deployment succeeds. An additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "002f438ecfea"
+      },
+      "outputs": [],
+      "source": [
+        "instances = [\n",
+        "    {\n",
+        "        \"text\": [\"a diagram\", \"a dog\", \"a cat\"],\n",
+        "        \"image\": image_to_base64(img_diagram),\n",
+        "    },\n",
+        "    {\n",
+        "        \"text\": [\"a diagram\", \"a dog\", \"two cats\", \"calico cat\"],\n",
+        "        \"image\": image_to_base64(img_cat),\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances).predictions\n",
+        "response"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "Clean up resources:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4fb87c6aeecf"
+      },
+      "source": [
+        "#### Image/text feature embedding"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf55e38815dc"
+      },
+      "outputs": [],
+      "source": [
+        "model, endpoint = deploy_model(\n",
+        "    checkpoint=f\"gs://{GCS_BUCKET}/{job_name}/checkpoints/epoch_latest.pt\",\n",
+        "    model=model_name,\n",
+        "    precision=precision,\n",
+        "    task=\"feature-embedding\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "80b3fd2ace09"
+      },
+      "source": [
+        "NOTE: The model weights are downloaded after the deployment succeeds. An additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "56d5c001575d"
+      },
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"text\": [\"a diagram\", \"a dog\", \"a cat\"],\n",
+        "        \"image\": image_to_base64(img_diagram),\n",
+        "    },\n",
+        "    {\n",
+        "        \"image\": image_to_base64(img_cat),\n",
+        "    },\n",
+        "    {\n",
+        "        \"text\": [\"a diagram\", \"a dog\", \"two cats\", \"calico cat\"],\n",
+        "    },\n",
+        "    {\n",
+        "        \"text\": \"a single value\",\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances).predictions\n",
+        "for pred in response:\n",
+        "    for k, v in pred.items():\n",
+        "        print(k, np.array(v).shape)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "Clean up resources:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_open_clip.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb

@@ -0,0 +1,887 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - OpenLLaMA (PEFT)\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a> (A Python-3 GPU notebook is recommended)\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates running local inference with prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm), finetuning and deploying OpenLLaMA with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)), and evaluating PEFT-finetuned OpenLLaMA in Vertex AI.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Run local inference with prebuilt OpenLLaMA\n",
+        "- Deploy prebuilt OpenLLaMA\n",
+        "- Deploy prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
+        "- Finetune and deploy OpenLLaMA with PEFT\n",
+        "- Evaluate finetuned OpenLLaMA with PEFT\n",
+        "\n",
+        "| Models | LoRA |\n",
+        "| :- | :- |\n",
+        "| [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b) | Y |\n",
+        "| [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) | Y |\n",
+        "| [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) | Y |\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/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": "264c07757582"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
+        "\n",
+        "Running local inference with OpenLLaMA requires a GPU."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ioensNKM8ned"
+      },
+      "source": [
+        "### Colab only\n",
+        "Run the following commands for Colab and skip this section if you are using Workbench."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "    # Install gdown for downloading example training images.\n",
+        "    ! pip3 install gdown\n",
+        "\n",
+        "    # Restart the notebook kernel after installs.\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "! pip3 install transformers==4.31.0\n",
+        "! pip3 install sentencepiece==0.1.99\n",
+        "! pip3 install accelerate==0.21.0"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bb7adab99e41"
+      },
+      "source": [
+        "### Setup Google Cloud project\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c460088b873"
+      },
+      "source": [
+        "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# Cloud project id.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The region you want to launch jobs in.\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The Cloud Storage bucket for storing experiments output.\n",
+        "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
+        "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
+        "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
+        "\n",
+        "# The service account looks like:\n",
+        "# '@.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "# The service account for deploying fine tuned model.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e828eb320337"
+      },
+      "source": [
+        "### Initialize Vertex AI API"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "12cd25839741"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built training and serving docker images.\n",
+        "TRAIN_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        ")\n",
+        "PREDICTION_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        ")\n",
+        "VLLM_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve\"\n",
+        ")\n",
+        "\n",
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "354da31189dc"
+      },
+      "outputs": [],
+      "source": [
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "\n",
+        "def create_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Creates a name with date time when triggering training or deployment\n",
+        "    jobs in Vertex AI.\n",
+        "    \"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"n1-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count: int = 1,\n",
+        ") -> tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    vllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=4\",\n",
+        "        \"--gpu-memory-utilization=0.95\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "65eaa62632d1"
+      },
+      "source": [
+        "## Run inferences locally with prebuilt OpenLLaMA"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "339601a9500b"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "from transformers import LlamaForCausalLM, LlamaTokenizer\n",
+        "\n",
+        "model_path = \"openlm-research/open_llama_3b\"\n",
+        "\n",
+        "tokenizer = LlamaTokenizer.from_pretrained(model_path)\n",
+        "\n",
+        "model = LlamaForCausalLM.from_pretrained(\n",
+        "    model_path,\n",
+        "    torch_dtype=torch.float16,\n",
+        "    device_map=\"auto\",\n",
+        ")\n",
+        "\n",
+        "prompt = \"Q: What is the largest animal?\\nA:\"\n",
+        "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
+        "input_ids = input_ids.to(\"cuda\")\n",
+        "generation_output = model.generate(input_ids=input_ids, max_new_tokens=32)\n",
+        "print(tokenizer.decode(generation_output[0]))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "V7VOhhHGpUrj"
+      },
+      "source": [
+        "## Deploy Prebuilt OpenLLaMA with vLLM\n",
+        "\n",
+        "This section deploys prebuilt OpenLLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "vLLM is a highly optimized LLM serving framework that can significantly increase serving throughput. The higher QPS you have, the more performance benefits you get from using vLLM."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4GTNnnuYqrW_"
+      },
+      "source": [
+        "Set the prebuilt model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "kLsRoc4Kqrkx"
+      },
+      "outputs": [],
+      "source": [
+        "prebuilt_model_id = \"openlm-research/open_llama_7b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YI0vaDi6p2fi"
+      },
+      "outputs": [],
+      "source": [
+        "model_without_peft_vllm, endpoint_without_peft_vllm = deploy_model_vllm(\n",
+        "    model_name=create_name_with_datetime(prefix=\"openllama-serve-vllm\"),\n",
+        "    model_id=prebuilt_model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=\"n1-highmem-8\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count=2,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dWYmYWoqqBuZ"
+      },
+      "source": [
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts. If you are interested in additional serving parameters, please refer to the vLLM GitHub [examples/api_client.py](https://github.com/vllm-project/vllm/blob/main/examples/api_client.py) for more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fjO4z3qAp3pK"
+      },
+      "outputs": [],
+      "source": [
+        "instance = {\n",
+        "    \"prompt\": \"Hi, Google.\",\n",
+        "    \"n\": 1,\n",
+        "    \"max_tokens\": 32,\n",
+        "}\n",
+        "response = endpoint_without_peft_vllm.predict(instances=[instance])\n",
+        "print(response.predictions[0])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e70e3519ff8b"
+      },
+      "source": [
+        "## Finetune and deploy OpenLLaMA with PEFT\n",
+        "\n",
+        "This section demonstrates how to finetune the OpenLLaMA-7b model, merge the finetuned LoRA adapter with the base model, and serve using vLLM."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5qCrm_kJH5cz"
+      },
+      "source": [
+        "Set the base model id."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "N3UBLiYrM3sU"
+      },
+      "outputs": [],
+      "source": [
+        "# vLLM currently does not support finetuned `open_llama_3b` model yet.\n",
+        "base_model_id = \"openlm-research/open_llama_7b\"  # @param [\"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iWGwJHqI7LMs"
+      },
+      "source": [
+        "### Finetune"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KKEYoRfiHDVv"
+      },
+      "source": [
+        "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
+        "\n",
+        "This example uses the dataset [Abirate/english_quotes](https://huggingface.co/datasets/Abirate/english_quotes).\n",
+        "\n",
+        "In order to make the finetuning efficiently, we enabled quantization for loading pretrained models for finetuning LoRA models. Precision options include `\"4bit\"`, `\"8bit\"`, `\"float16\"` (default) and `\"float32\"`, and the precision can be set via `\"--precision_mode\"`. The peak GPU memory usages are ~7G, ~10G and ~16G for finetuning LoRA models for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) separately with default training parameters and the example dataset. `open_llama_3b` and `open_llama_7b` can be finetuned on **1 V100**, and `open_llama_13b` can be finetuned on **1 A100 (40G)**.\n",
+        "\n",
+        "In this section, the finetuned LoRA adapter will be saved to a GCS bucket specified by the variable `lora_adapter_dir` below; and we merge the LoRa adapter with the base model, and save it to a separate GCS bucket specified by `merged_model_output_dir` below.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "65467b361315"
+      },
+      "outputs": [],
+      "source": [
+        "dataset_name = \"Abirate/english_quotes\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Finetunes open_llama_7b with 1 V100.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "# Finetunes open_llama_13b with 1 A100 (40G).\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "replica_count = 1\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Setup training job.\n",
+        "job_name = create_name_with_datetime(\"openllama-lora-train\")\n",
+        "train_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "\n",
+        "# Create a GCS folder to store the LORA adapter.\n",
+        "lora_adapter_dir = create_name_with_datetime(\"openllama-lora-adapter\")\n",
+        "lora_output_dir = os.path.join(MODEL_BUCKET, lora_adapter_dir)\n",
+        "lora_output_dir_gcsfuse = lora_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Create a GCS folder to store the merged model with the base model and the\n",
+        "# finetuned LORA adapter.\n",
+        "merged_model_dir = create_name_with_datetime(\"openllama-merged-model\")\n",
+        "merged_model_output_dir = os.path.join(MODEL_BUCKET, merged_model_dir)\n",
+        "merged_model_output_dir_gcsfuse = merged_model_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Pass training arguments and launch job.\n",
+        "train_job.run(\n",
+        "    args=[\n",
+        "        \"--task=causal-language-modeling-lora\",\n",
+        "        f\"--pretrained_model_id={base_model_id}\",\n",
+        "        f\"--dataset_name={dataset_name}\",\n",
+        "        f\"--output_dir={lora_output_dir_gcsfuse}\",\n",
+        "        f\"--merge_base_and_lora_output_dir={merged_model_output_dir_gcsfuse}\",\n",
+        "        \"--lora_rank=16\",\n",
+        "        \"--lora_alpha=32\",\n",
+        "        \"--lora_dropout=0.05\",\n",
+        "        \"--warmup_steps=10\",\n",
+        "        \"--max_steps=10\",\n",
+        "        \"--learning_rate=2e-4\",\n",
+        "    ],\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    boot_disk_size_gb=500,\n",
+        ")\n",
+        "\n",
+        "print(\"The finetuned Lora adapter can be found at: \", lora_output_dir)\n",
+        "print(\n",
+        "    \"The finetuned Lora adapter merged with the base model can be found at: \",\n",
+        "    merged_model_output_dir,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jqmCtkGnhDmp"
+      },
+      "source": [
+        "### Deploy with vLLM\n",
+        "This section uploads the model to Model Registry and deploys it on the Endpoint. vLLM currently does not support serving finetuned [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b) for now so we will use the [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) in this example.\n",
+        "\n",
+        "The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "The peak GPU memory usages for [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) with LoRA weights are ~8.7G and ~15.2G respectively with the default settings.\n",
+        "\n",
+        "NOTE: vLLM requires a merged model with the base model and the finetuned LoRA adapter. Based on your business need, if you need the base model and the finetuned LoRA weight to be served separately, please consider using the regular Vertex serving instead.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf55e38815dc"
+      },
+      "outputs": [],
+      "source": [
+        "model_with_peft_vllm, endpoint_with_peft_vllm = deploy_model_vllm(\n",
+        "    model_name=create_name_with_datetime(prefix=\"openllama-peft-serve-vllm\"),\n",
+        "    model_id=merged_model_output_dir,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=\"n1-highmem-8\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count=2,\n",
+        ")\n",
+        "\n",
+        "print(\"endpoint_name:\", endpoint_with_peft_vllm.name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "80b3fd2ace09"
+      },
+      "source": [
+        "NOTE: After the deployment succeeds, the base model weights will be downloaded on the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4ab04da3ec9a"
+      },
+      "outputs": [],
+      "source": [
+        "instance = {\n",
+        "    \"prompt\": \"Hi, Google. How are you doing?\",\n",
+        "    \"n\": 1,\n",
+        "    \"max_tokens\": 32,\n",
+        "}\n",
+        "response = endpoint_with_peft_vllm.predict(instances=[instance])\n",
+        "print(response.predictions[0])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vhDf9dNNn4bP"
+      },
+      "source": [
+        "### [Optional] Merge a previously trained LoRA adapter with the base model\n",
+        "\n",
+        "This section demonstrates how to merge a previously trained LoRA adapter with a base model, and save the merged model to a GCS bucket. Please be aware that the LoRA adapter should be trained on the same base model.\n",
+        "\n",
+        "In this example, we use the [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) model for the demonstartion."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YHdru1aRqRFF"
+      },
+      "outputs": [],
+      "source": [
+        "merge_job_name = create_name_with_datetime(prefix=\"openllama-peft-merge\")\n",
+        "\n",
+        "# The base model to be merged upon. It can be a huggingface model id, or a GCS\n",
+        "# path where the base model was stored.\n",
+        "base_model_dir = \"gs://\"  # @param {type:\"string\"}\n",
+        "# The previously trained LoRA adapter. It needs to be stored in a GCS path.\n",
+        "finetuned_lora_adapter_dir = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The GCS path to save the merged model\n",
+        "merged_model_output_dir = os.path.join(MODEL_BUCKET, merge_job_name)\n",
+        "merged_model_output_dir_gcsfuse = merged_model_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Merge base model `open_llama_7b` with 1 V100.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "# Merge base model `open_llama_13b` with 1 A100 (40G).\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": machine_type,\n",
+        "            \"accelerator_type\": accelerator_type,\n",
+        "            \"accelerator_count\": 1,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAIN_DOCKER_URI,\n",
+        "            \"command\": [],\n",
+        "            \"args\": [\n",
+        "                \"--task=merge-causal-language-model-lora\",\n",
+        "                \"--merge_model_precision_mode=float16\",\n",
+        "                \"--pretrained_model_id=%s\" % base_model_dir,\n",
+        "                \"--finetuned_lora_model_dir=%s\" % finetuned_lora_adapter_dir,\n",
+        "                \"--merge_base_and_lora_output_dir=%s\" % merged_model_output_dir_gcsfuse,\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "merge_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=merge_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "merge_custom_job.run()\n",
+        "\n",
+        "print(\"The merged model is stored at: \", merged_model_output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JmuUk3l1DoEo"
+      },
+      "source": [
+        "## Evaluate PEFT-finetuned OpenLLaMA\n",
+        "\n",
+        "This section demonstrates how to evaluate the OpenLLaMA model fintuned with PEFT LoRA using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob.\n",
+        "\n",
+        "This example uses the dataset [HellaSwag](https://allenai.org/data/hellaswag). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gM4SXaquDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "eval_dataset = \"hellaswag\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Evaluates open_llama_3b and open_llama_7b with 1 V100.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "# Evaluates open_llama_13b with 1 A100 (40G).\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "replica_count = 1\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "job_name = create_name_with_datetime(prefix=\"openllama-peft-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Yt4Tth8hDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "# Prepare evaluation script that runs the evaluation harness.\n",
+        "script_path = \"./eval_script.py\"  # @param {type:\"string\"}\n",
+        "\n",
+        "eval_command = f\"\"\"import subprocess\n",
+        "\n",
+        "\n",
+        "subprocess.call([\n",
+        "    'python',\n",
+        "    'main.py',\n",
+        "    '--model',\n",
+        "    'hf-causal-experimental',\n",
+        "    '--model_args',\n",
+        "    'pretrained={base_model_id},peft={output_dir_gcsfuse}',\n",
+        "    '--tasks',\n",
+        "    '{eval_dataset}',\n",
+        "    '--output_path',\n",
+        "    '{eval_output_dir_gcsfuse}',\n",
+        "])\n",
+        "\"\"\"\n",
+        "\n",
+        "with open(script_path, \"w\") as fp:\n",
+        "    fp.write(eval_command)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ItWB0WS__CX-"
+      },
+      "source": [
+        "### Submit evaluation CustomJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BbyIk99bDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "# Pass evaluation arguments and launch job.\n",
+        "eval_job = aiplatform.CustomJob.from_local_script(\n",
+        "    display_name=job_name,\n",
+        "    script_path=script_path,\n",
+        "    container_uri=EVAL_DOCKER_URI,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    base_output_dir=eval_output_dir,\n",
+        ")\n",
+        "\n",
+        "eval_job.run()\n",
+        "\n",
+        "print(\"Evaluation results were saved in:\", eval_output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kN0lE2iu_NXN"
+      },
+      "source": [
+        "### Fetch and print evaluation results"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "927oRxoADoEp"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "from google.cloud import storage\n",
+        "\n",
+        "# Fetch evaluation results.\n",
+        "storage_client = storage.Client()\n",
+        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
+        "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
+        "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :]\n",
+        "blob = bucket.blob(RESULT_FILE_PATH)\n",
+        "raw_result = blob.download_as_string()\n",
+        "\n",
+        "# Print evaluation results.\n",
+        "result = json.loads(raw_result)\n",
+        "result_formatted = json.dumps(result, indent=2)\n",
+        "print(f\"Evaluation result:\\n{result_formatted}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete custom train and evaluation jobs.\n",
+        "train_job.delete()\n",
+        "eval_job.delete()\n",
+        "\n",
+        "# Undeploy models and delete endpoints.\n",
+        "endpoint_without_peft_vllm.delete(force=True)\n",
+        "endpoint_with_peft_vllm.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_without_peft_vllm.delete()\n",
+        "model_with_peft_vllm.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_openllama_peft.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_owlvit.ipynb

@@ -204,7 +204,7 @@
       "outputs": [],
       "source": [
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-transformers-serve\""
       ]
     },
     {

notebooks/community/model_garden/model_garden_pytorch_peft.ipynb

@@ -246,8 +246,8 @@
       "outputs": [],
       "source": [
         "# The pre-built training and serving docker images.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
-        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-serve\""
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\""
       ]
     },
     {
@@ -649,14 +649,13 @@
         "# # If deploy finetuned falcon-40b-instruct models, please set\n",
         "# machine_type = \"a2-highgpu-1g\",\n",
         "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
-        "machine_type = \"n1-standard-8\",\n",
+        "machine_type = \"n1-standard-8\"\n",
         "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
         "\n",
-        "accelerator_type =\n",
         "model, endpoint = deploy_model(\n",
         "    model_name=get_job_name_with_datetime(prefix=\"falcon-peft-serve\"),\n",
         "    base_model_id=base_model_id,\n",
-        "    finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\"+str(max_steps)),\n",
+        "    finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\" + str(max_steps)),\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    task=\"instruct-lora\",\n",
         "    machine_type=machine_type,\n",

notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb

@@ -0,0 +1,633 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "ad41232f-4ac9-4607-8ee6-9b1d98d0f5c3",
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2023 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."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "30e2cb63",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Pic2Word\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pic2word.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jpic2word.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>                                                                                               <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pic2word.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "5d5af88c-a3c8-46fa-b90a-1b6737a91534",
+      "metadata": {
+        "id": "7e3e5205fbfd"
+      },
+      "source": [
+        "## Overview"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "87b71198-5c24-4e8c-aab1-44e080aa7fe7",
+      "metadata": {
+        "id": "523084ce4894"
+      },
+      "source": [
+        "This notebook demonstrates how to use the [Pic2Word](https://github.com/google-research/composed_image_retrieval) model in Vertex AI Model Garden. "
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "0ba32ec8-e06f-474f-b753-20aaf74ebd12",
+      "metadata": {
+        "id": "f9cbeb1704e1"
+      },
+      "source": [
+        "## Objective"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "95694bc4-0626-4371-8b2a-6b89967742db",
+      "metadata": {
+        "id": "da71cd41e69e"
+      },
+      "source": [
+        "Following the notebook you will conduct experiments using the pre-built docker image on Vertex AI.\n",
+        "\n",
+        "- Deploy pretrained Pic2Word models in Google Cloud Vertex AI\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI Online Prediction"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "b1893e0c-859b-4d42-b1d8-d133855d8767",
+      "metadata": {
+        "id": "8264c04e0f34"
+      },
+      "source": [
+        "## Dataset"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "a099439c-3122-429b-8215-caaf9aee69c5",
+      "metadata": {
+        "id": "52e4f9df87b2"
+      },
+      "source": [
+        "We use the [COCO](https://cocodataset.org/#home) validation set (5,000 images) for evaluation."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "56d2a9f8-2d99-407e-a378-57f99754185b",
+      "metadata": {
+        "id": "169ef74e43a6"
+      },
+      "source": [
+        "## Costs"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "bb62de0d-71ac-45ee-82ac-5f7465f2f0eb",
+      "metadata": {
+        "id": "61a64432ac87"
+      },
+      "source": [
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "- Vertex AI\n",
+        "- Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing?_ga=2.46650789.-341051769.1686949237) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing?_ga=2.46650789.-341051769.1686949237), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/?_ga=2.247379078.-341051769.1686949237) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "ea27403e",
+      "metadata": {
+        "id": "586c9147ba8a"
+      },
+      "source": [
+        "# Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "f9c685d9",
+      "metadata": {
+        "id": "68f1241e68a7"
+      },
+      "outputs": [],
+      "source": [
+        "if \"google.colab\" in str(get_ipython()):\n",
+        "    # Configs for colab notebooks.\n",
+        "    ! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "e75fe500-5dc1-4ab3-8829-3eeada5be7f7",
+      "metadata": {
+        "id": "d0c7bca32b78"
+      },
+      "source": [
+        "## Setup environment"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "e8b0ef65",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with **Vertex AI User** and **Storage Object Admin** roles for deploying fine tuned model to Vertex AI endpoint.\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
+        "\n",
+        "**Note**: Jupyter runs lines prefixed with ! as shell commands, and it interpolates Python variables prefixed with $ into these commands."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "09490824",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "6ddc1d11",
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "9ed7d668",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "90d0108f",
+      "metadata": {
+        "id": "twgKk-LsLmX3"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "ca43938d",
+      "metadata": {
+        "id": "409f1effab6c"
+      },
+      "source": [
+        "### Buckets\n",
+        "\n",
+        "You can create a storage bucket to store model input and output images."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "889d10ac",
+      "metadata": {
+        "id": "bcdadd216c34"
+      },
+      "outputs": [],
+      "source": [
+        "# The form for BUCKET_URI is gs://.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "INPUT_BUCKET = os.path.join(BUCKET_URI, \"/input\")\n",
+        "OUTPUT_BUCKET = os.path.join(BUCKET_URI, \"/output\")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "15c0e671",
+      "metadata": {
+        "id": "sBCra4QMA2wR"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "623a03cc",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "fbee4160",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "3eabaf6b",
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "90a615fe",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "3d4e50e1",
+      "metadata": {
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "b0b40233",
+      "metadata": {
+        "id": "f6b2ccc891ed"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "a8019fcc-26ce-4536-b394-aa1fb18794ee",
+      "metadata": {
+        "id": "4226467373ef"
+      },
+      "source": [
+        "If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk) and [gsutil](https://cloud.google.com/storage/docs/gsutil_install)."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "315598c1-3ff4-4bc8-9b82-7e2fe55d4678",
+      "metadata": {
+        "id": "8ba36d3a171d"
+      },
+      "source": [
+        "### Setup variables"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "6d26ddb2-04ad-4973-8965-9b25eb24993f",
+      "metadata": {
+        "id": "25ca675e59e7"
+      },
+      "outputs": [],
+      "source": [
+        "# Prediction constants.\n",
+        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pic2word_serve:latest\"\n",
+        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
+        "\n",
+        "# The serving port.\n",
+        "SERVE_PORT = 7080\n",
+        "\n",
+        "# The service account looks like:\n",
+        "# '@.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "0daec6e5-7642-458f-bd93-60a15431b0a1",
+      "metadata": {
+        "id": "71dd15118703"
+      },
+      "source": [
+        "## Deploy model for online prediction"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "b286ce6b-cb00-441b-8249-93497de17974",
+      "metadata": {
+        "id": "5c891f6352ad"
+      },
+      "source": [
+        "This section uploads the model to Vertex Model Registry and deploys it on an Endpoint resource. This will take around 15 minutes to finish."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "01f686b2-9227-496b-9c4d-a63becb59fb8",
+      "metadata": {
+        "id": "e6abf4ee450a"
+      },
+      "source": [
+        "### Upload and deploy model to Vertex AI"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "8885d5be-5007-4778-9562-976102fefe73",
+      "metadata": {
+        "id": "645d8d1df8d4"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Init common setup.\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
+        "\n",
+        "# Upload model.\n",
+        "serving_env = {}\n",
+        "model = aiplatform.Model.upload(\n",
+        "    display_name=\"pic2word-model\",\n",
+        "    serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
+        "    serving_container_ports=[SERVE_PORT],\n",
+        "    serving_container_predict_route=\"/predictions/pic2word\",\n",
+        "    serving_container_health_route=\"/ping\",\n",
+        "    serving_container_environment_variables=serving_env,\n",
+        ")\n",
+        "# Or reuse a pre-uploaded model.\n",
+        "# model = aiplatform.Model('projects/123456789/locations/us-central1/models/123456789@1')\n",
+        "\n",
+        "# Create an endpoint.\n",
+        "endpoint = aiplatform.Endpoint.create(display_name=\"pytorch-pic2word-endpoint\")\n",
+        "# Or reuse a pre-created endpoint.\n",
+        "# endpoint = aiplatform.Endpoint('projects/123456789/locations/us-central1/endpoints/123456789')\n",
+        "\n",
+        "# Deploy model to endpoint.\n",
+        "model.deploy(\n",
+        "    endpoint=endpoint,\n",
+        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "    accelerator_count=1,\n",
+        "    traffic_percentage=100,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        ")"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "f0bd36bb-f988-479a-be40-18c54980ca1f",
+      "metadata": {
+        "id": "2bb47a2855a2"
+      },
+      "source": [
+        "You can manage your uploaded models in the [Model Registry](https://pantheon.corp.google.com/vertex-ai/models) and your endpoints in the [Endpoints](https://pantheon.corp.google.com/vertex-ai/endpoints)."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "95edbcdd-8246-444a-aef2-a9cadbc03000",
+      "metadata": {
+        "id": "f14ff5e7a476"
+      },
+      "source": [
+        "## Send a prediction request to the endpoint"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "3a83d6bb",
+      "metadata": {
+        "id": "c41c7d51d654"
+      },
+      "source": [
+        "Set the query text and upload your desired query image to the specified INPUT_BUCKET. Use an asterisk in the query text to specify which token gets replaced by the image token. "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "f7c6fc99-42cd-4544-86df-b14f49189bfd",
+      "metadata": {
+        "id": "58d4f8e84e02"
+      },
+      "outputs": [],
+      "source": [
+        "payload = json.dumps(\n",
+        "    {\n",
+        "        \"query\": \"a bunch of *\",\n",
+        "        \"image_path\": INPUT_BUCKET,\n",
+        "        \"image_file_name\": \"cat.png\",\n",
+        "        \"output_storage_dir\": OUTPUT_BUCKET,\n",
+        "    }\n",
+        ")\n",
+        "response = endpoint.predict(payload).predictions\n",
+        "print(response)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "id": "8e71ee23-967c-4871-802f-ccb39958386e",
+      "metadata": {
+        "id": "608b936179df"
+      },
+      "source": [
+        "## Cleaning Up\n",
+        "\n",
+        "To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "d9be2a18-db11-4108-9622-647b33be2594",
+      "metadata": {
+        "id": "7d17b385141d"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete endpoint resource.\n",
+        "endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete model resource.\n",
+        "model.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created.\n",
+        "delete_bucket = False\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_pic2word.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb

@@ -24,7 +24,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "99c1c3fc2ca5"
@@ -63,14 +62,16 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This notebook demonstrates running local inference for [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) on either [Colab](https://colab.research.google.com) or [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench). This notebook also demonstrates finetuning runwayml/stable-diffusion-v1-5 with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth) and deploying it on Vertex AI for online prediction.\n",
+        "This notebook demonstrates running local inference for [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) on either [Colab](https://colab.research.google.com) or [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench). This notebook also demonstrates finetuning runwayml/stable-diffusion-v1-5 with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth), finetuning with [LoRA](https://huggingface.co/docs/diffusers/training/text2image#lora), and finetuning with [Dreambooth and LoRA](https://huggingface.co/docs/diffusers/training/dreambooth#finetuning-with-lora) and deploying it on Vertex AI for online prediction.\n",
         "\n",
         "### Objective\n",
         "\n",
         "- Run local predictions for text-to-image and text-guided-image-to-image with serving dockers.\n",
-        "- Finetune the stable-diffusion-v1.5 model with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth).\n",
-        "- Upload the model to [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
-        "- Deploy the model to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
+        "- Full parameter finetune stable-diffusion-v1.5 model with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth).\n",
+        "- Finetune the stable-diffusion-v1.5 model with [LoRA](https://huggingface.co/docs/diffusers/training/text2image#lora).\n",
+        "- Finetune the stable-diffusion-v1.5 model with [Dreambooth + LoRA](https://huggingface.co/docs/diffusers/training/dreambooth#finetuning-with-lora).\n",
+        "- Upload models to [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
+        "- Deploy models to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
         "- Run online predictions for text-to-image and text-guided-image-to-image.\n",
         "\n",
         "### Costs\n",
@@ -84,7 +85,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "264c07757582"
@@ -96,7 +96,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ioensNKM8ned"
@@ -106,7 +105,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "d73ffa0c0b83"
@@ -132,9 +130,9 @@
         "    ! pip3 install --upgrade pip\n",
         "    ! pip3 install torchvision==0.14.1\n",
         "    ! pip3 install transformers==4.27.1\n",
-        "    ! pip3 install diffusers==0.15.1\n",
+        "    ! pip3 install diffusers==0.20.1\n",
         "    ! pip3 install datasets==2.9.0\n",
-        "    ! pip3 install accelerate==0.18.0\n",
+        "    ! pip3 install accelerate==0.21.0\n",
         "    ! pip3 install triton==2.0.0.dev20221120\n",
         "    ! pip3 install xformers==0.0.16\n",
         "    # Install gdown for downloading example training images.\n",
@@ -150,7 +148,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "fb671e75ca7b"
@@ -167,7 +164,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "bb7adab99e41"
@@ -187,7 +183,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "6c460088b873"
@@ -218,7 +213,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "e828eb320337"
@@ -241,7 +235,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "2cc825514deb"
@@ -258,15 +251,20 @@
       },
       "outputs": [],
       "source": [
-        "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "# The pre-built training docker images. They contain training scripts and models.\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "PEFT_TRAIN_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        ")\n",
         "\n",
-        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "# The pre-built serving docker images. They contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\"\n",
+        "PEFT_SERVE_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        ")"
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "0c250872074f"
@@ -352,6 +350,41 @@
         "    return model, endpoint\n",
         "\n",
         "\n",
+        "def deploy_model_peft(\n",
+        "    model_name,\n",
+        "    base_model_id,\n",
+        "    finetuned_lora_model_path,\n",
+        "    task,\n",
+        "    machine_type=\"n1-standard-8\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "):\n",
+        "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    serving_env = {\n",
+        "        \"BASE_MODEL_ID\": base_model_id,\n",
+        "        \"TASK\": task,\n",
+        "    }\n",
+        "    if finetuned_lora_model_path:\n",
+        "        serving_env[\"FINETUNED_LORA_MODEL_PATH\"] = finetuned_lora_model_path\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=PEFT_SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predictions/peft_serving\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=1,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
         "def get_bucket_and_blob_name(filepath):\n",
         "    # The gcs path is of the form gs://<bucket-name>/<blob-name>\n",
         "    gs_suffix = filepath.split(\"gs://\", 1)[1]\n",
@@ -375,7 +408,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "e8a42fa49305"
@@ -385,7 +417,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "1d5ebc91c786"
@@ -419,7 +450,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "5aed5ed7b6f6"
@@ -462,7 +492,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "e70e3519ff8b"
@@ -472,17 +501,18 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "0dc65d8f0689"
+        "id": "yIIgB6Fw_6Qk"
       },
       "source": [
+        "### Finetune\n",
+        "\n",
         "This section uses [dreambooth](https://dreambooth.github.io/) to finetune the [stable-diffusion-v1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5) model with [5 dog images](https://drive.google.com/drive/folders/1BO_dyz-p65qhBRRMRA4TbZ8qW4rB99JZ) to personalize the text-to-image model.\n",
         "\n",
         "It finetunes both text encoder and unet of the stable diffusion model up to 800 steps. The whole finetuning job takes 30 minutes to finish using 1 A100 GPU.\n",
         "\n",
-        "The full model will be saved after the finetuning job finishs and it can be loaded by the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) to run inference."
+        "The full model will be saved after the finetuning job finishes and it can be loaded by the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) to run inference."
       ]
     },
     {
@@ -502,7 +532,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "969cfeb79317"
@@ -541,7 +570,7 @@
         ")\n",
         "\n",
         "# Pass training arguments and launch job.\n",
-        "# See https://github.com/huggingface/diffusers/blob/v0.14.0/examples/dreambooth/train_dreambooth.py#L75\n",
+        "# See https://github.com/huggingface/diffusers/blob/v0.16.0/examples/dreambooth/train_dreambooth.py#L75\n",
         "# for a full list of training arguments.\n",
         "model = job.run(\n",
         "    args=[\n",
@@ -572,39 +601,28 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "bf7f82732e61"
       },
       "source": [
-        "## Upload and Deploy models"
-      ]
-    },
-    {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1cc26e68d7b0"
-      },
-      "source": [
+        "### Upload and Deploy models\n",
+        "\n",
         "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
         "\n",
         "The model deployment step will take ~15 minutes to complete."
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "cd7b56421392"
       },
       "source": [
-        "### Text-to-image"
+        "#### Text-to-image"
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "6d331b1ea337"
@@ -626,13 +644,12 @@
       "outputs": [],
       "source": [
         "# Set the model_id to \"runwayml/stable-diffusion-v1-5\" to load the OSS pre-trained model.\n",
-        "model, endpoint = deploy_model(\n",
+        "model_text_to_image, endpoint_text_to_image = deploy_model(\n",
         "    model_id=f\"gs://{GCS_BUCKET}/dreambooth/output\", task=\"text-to-image\"\n",
         ")"
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "80b3fd2ace09"
@@ -650,53 +667,26 @@
       "outputs": [],
       "source": [
         "instances = [\n",
-        "    {\"prompt\": \"a squirrel in Picasso style\"},\n",
-        "    {\"prompt\": \"a dog in Picasso style\"},\n",
-        "    {\"prompt\": \"a cat in Picasso style\"},\n",
-        "    {\"prompt\": \"a deer in Picasso style\"},\n",
+        "    {\"prompt\": \"A picture of a sks dog in a house\"},\n",
+        "    {\"prompt\": \"A picture of a sks dog catching a frisbee\"},\n",
+        "    {\"prompt\": \"A picture of a sks dog in front of a computer\"},\n",
+        "    {\"prompt\": \"A picture of a sks dog in a bucket\"},\n",
         "]\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoint_text_to_image.predict(instances=instances)\n",
         "images = [base64_to_image(image) for image in response.predictions]\n",
         "image_grid(images)"
       ]
     },
     {
-      "attachments": {},
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "af21a3cff1e0"
-      },
-      "source": [
-        "Clean up resources:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "911406c1561e"
-      },
-      "outputs": [],
-      "source": [
-        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
-        "\n",
-        "# Delete models.\n",
-        "model.delete()"
-      ]
-    },
-    {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "c1e51f764a60"
       },
       "source": [
-        "### Text-guided image-to-image"
+        "#### Text-guided image-to-image"
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "fa686a54047c"
@@ -714,13 +704,12 @@
       "outputs": [],
       "source": [
         "# Set the model_id to a GCS path, like \"gs://GCS_BUCKET/dreambooth/output\", to load the dreambooth finetuned model above.\n",
-        "model, endpoint = deploy_model(\n",
+        "model_image_to_image, endpoint_image_to_image = deploy_model(\n",
         "    model_id=\"runwayml/stable-diffusion-v1-5\", task=\"image-to-image\"\n",
         ")"
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "SsjYFLvNymc0"
@@ -747,19 +736,326 @@
         "        \"image\": image_to_base64(init_image),\n",
         "    },\n",
         "]\n",
-        "response = endpoint.predict(instances=instances)\n",
+        "response = endpoint_image_to_image.predict(instances=instances)\n",
         "images = [base64_to_image(image) for image in response.predictions]\n",
         "display(images[0])"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Z3-T96nh1Q8K"
+      },
+      "source": [
+        "## Finetune and deploy with LoRA"
+      ]
+    },
     {
       "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EY-iQTub1UJg"
+      },
+      "source": [
+        "### Finetune\n",
+        "\n",
+        "This section uses [LoRA](https://arxiv.org/abs/2106.09685) to finetune the [stable-diffusion-v1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5) model with [lambdalabs/pokemon-blip-captions](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions).\n",
+        "\n",
+        "Finetuning with LoRA \n",
+        "\n",
+        "The LoRA weights will be saved after the finetuning job finishes and it can be loaded by the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) to run inference."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "J6Q5CMgX2py9"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-trained model to be loaded.\n",
+        "model_id = \"runwayml/stable-diffusion-v1-5\"\n",
+        "DATASET_NAME = \"lambdalabs/pokemon-blip-captions\"\n",
+        "\n",
+        "# Output path.\n",
+        "output_dir_lora = f\"/gcs/{GCS_BUCKET}/lora/output\"\n",
+        "gs_output_dir_lora = f\"gs://{GCS_BUCKET}/lora/output\"\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "num_nodes = 1\n",
+        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
+        "num_gpus = 1\n",
+        "\n",
+        "\n",
+        "# Setup training job.\n",
+        "job_name = create_job_name(\"lora-stable-diffusion\")\n",
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=PEFT_TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "\n",
+        "# Set task to \"text-to-image-lora\" to finetune with LoRA.\n",
+        "# Pass training arguments and launch job.\n",
+        "# See https://github.com/huggingface/diffusers/blob/87ae330056f6942817656c8f7146283e90cf986b/examples/text_to_image/train_text_to_image_lora.py#L84\n",
+        "# for a full list of training arguments.\n",
+        "model = job.run(\n",
+        "    args=[\n",
+        "        \"--task=text-to-image-lora\",\n",
+        "        f\"--pretrained_model_name_or_path={model_id}\",\n",
+        "        f\"--dataset_name={DATASET_NAME}\",\n",
+        "        \"--caption_column=text\",\n",
+        "        \"--resolution=512\",\n",
+        "        \"--random_flip\",\n",
+        "        \"--train_batch_size=1\",\n",
+        "        \"--num_train_epochs=100\",\n",
+        "        \"--checkpointing_steps=5000\",\n",
+        "        \"--learning_rate=1e-04\",\n",
+        "        \"--lr_scheduler=constant\",\n",
+        "        \"--lr_warmup_steps=0\",\n",
+        "        f\"--output_dir={output_dir_lora}\",\n",
+        "        \"--validation_prompt='cute dragon creature'\",\n",
+        "        \"--seed=42\",\n",
+        "    ],\n",
+        "    replica_count=num_nodes,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=gpu_type,\n",
+        "    accelerator_count=num_gpus,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "AairVI_a7yyr"
+      },
+      "source": [
+        "### Deploy\n",
+        "\n",
+        "Deploy the stable diffusion model for the text-to-image task.\n",
+        "\n",
+        "Once deployed, you can send a batch of text prompts to the endpoint to generated images."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e154LQnB8Pap"
+      },
+      "outputs": [],
+      "source": [
+        "model_with_peft, endpoint_with_peft = deploy_model_peft(\n",
+        "    model_name=create_job_name(prefix=\"lora-stable-diffusion\"),\n",
+        "    base_model_id=model_id,\n",
+        "    finetuned_lora_model_path=gs_output_dir_lora,\n",
+        "    task=\"text-to-image-lora\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "h1f019sH8kse"
+      },
+      "source": [
+        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed after the above model deployment step succeeds and before you run the next step below. Otherwise you might see a ServiceUnavailable: 503 502:Bad Gateway error when you send requests to the endpoint."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WVPCsaBp8odc"
+      },
+      "outputs": [],
+      "source": [
+        "instances = [\n",
+        "    {\"prompt\": \"A pokemon with green eyes and red legs.\"},\n",
+        "    {\"prompt\": \"A pokemon with yellow tails and blue arms.\"},\n",
+        "    {\"prompt\": \"A pokemon with three heads and green feet.\"},\n",
+        "    {\"prompt\": \"A dragon pokemon with red and white stripes.\"},\n",
+        "]\n",
+        "response = endpoint_with_peft.predict(instances=instances)\n",
+        "images = [base64_to_image(image) for image in response.predictions]\n",
+        "image_grid(images)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8jTetcepAZeO"
+      },
+      "source": [
+        "## Finetune with Dreambooth and LoRA"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9gQFXhVuAcc-"
+      },
+      "source": [
+        "### Finetune\n",
+        "\n",
+        "This section uses [dreambooth](https://dreambooth.github.io/) and [LoRA](https://arxiv.org/abs/2106.09685) to finetune the [stable-diffusion-v1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5) model with [5 dog images](https://drive.google.com/drive/folders/1BO_dyz-p65qhBRRMRA4TbZ8qW4rB99JZ) to personalize the text-to-image model.\n",
+        "\n",
+        "The full model will be saved after the finetuning job finishes and it can be loaded by the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) to run inference.\n",
+        "\n",
+        "Note: When using LoRA we can use a much higher learning rate compared to vanilla dreambooth. Here we use `1e-4` instead of the usual `2e-6`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dh1JNKYDUuwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# Download example training images.\n",
+        "!gdown --folder https://drive.google.com/drive/folders/1BO_dyz-p65qhBRRMRA4TbZ8qW4rB99JZ\n",
+        "\n",
+        "# Upload data to Cloud Storage bucket.\n",
+        "upload_local_dir_to_gcs(\"dog\", f\"gs://{GCS_BUCKET}/dreambooth-lora/dog\")\n",
+        "upload_local_dir_to_gcs(\"dog\", f\"gs://{GCS_BUCKET}/dreambooth-lora/dog_class\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "22zEZjQKU4R-"
+      },
+      "source": [
+        "**NOTE**: If the upload step fails due to lacking of permission, you need to [grant the Storage Object Admin role](https://cloud.google.com/storage/docs/access-control/using-iam-permissions) for the Cloud account of the notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "QKUXH9b9Bvta"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-trained model to be loaded.\n",
+        "model_id = \"runwayml/stable-diffusion-v1-5\"\n",
+        "\n",
+        "# Input and output path.\n",
+        "instance_dir = f\"/gcs/{GCS_BUCKET}/dreambooth-lora/dog\"\n",
+        "output_dir_dreambooth_lora = f\"/gcs/{GCS_BUCKET}/dreambooth-lora/output\"\n",
+        "gs_output_dir_dreambooth_lora = f\"gs://{GCS_BUCKET}/dreambooth-lora/output\"\n",
+        "\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "machine_type = \"n1-highgpu-8\"\n",
+        "num_nodes = 1\n",
+        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
+        "num_gpus = 1\n",
+        "\n",
+        "# Setup training job.\n",
+        "job_name = create_job_name(\"dreambooth-lora-stable-diffusion\")\n",
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=PEFT_TRAIN_DOCKER_URI,\n",
+        ")\n",
+        "\n",
+        "# Set task to \"text-to-image-dreambooth-lora\" to finetune using Dreambooth and\n",
+        "# LoRA.\n",
+        "# Pass training arguments and launch job.\n",
+        "# See https://github.com/huggingface/diffusers/blob/87ae330056f6942817656c8f7146283e90cf986b/examples/dreambooth/train_dreambooth_lora.py#L142\n",
+        "# for a full list of training arguments.\n",
+        "model = job.run(\n",
+        "    args=[\n",
+        "        \"--task=text-to-image-dreambooth-lora\",\n",
+        "        f\"--pretrained_model_name_or_path={model_id}\",\n",
+        "        f\"--instance_data_dir={instance_dir}\",\n",
+        "        f\"--output_dir={output_dir_dreambooth_lora}\",\n",
+        "        \"--instance_prompt='a photo of sks dog'\",\n",
+        "        \"--resolution=512\",\n",
+        "        \"--train_batch_size=1\",\n",
+        "        \"--gradient_accumulation_steps=1\",\n",
+        "        \"--checkpointing_steps=100\",\n",
+        "        \"--learning_rate=1e-4\",\n",
+        "        \"--lr_scheduler=constant\",\n",
+        "        \"--lr_warmup_steps=0\",\n",
+        "        \"--max_train_steps=500\",\n",
+        "        \"--validation_prompt='A photo of sks dog in a bucket'\",\n",
+        "        \"--validation_epochs=50\",\n",
+        "        \"--seed=0\",\n",
+        "    ],\n",
+        "    replica_count=num_nodes,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=gpu_type,\n",
+        "    accelerator_count=num_gpus,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iqrVgsVODmZb"
+      },
+      "source": [
+        "### Deploy\n",
+        "\n",
+        "Deploy the stable diffusion model for the text-to-image task.\n",
+        "\n",
+        "Once deployed, you can send a batch of text prompts to the endpoint to generated images."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ye3_N5bxDmZc"
+      },
+      "outputs": [],
+      "source": [
+        "model_with_peft_dreambooth, endpoint_with_peft_dreambooth = deploy_model_peft(\n",
+        "    model_name=create_job_name(prefix=\"dreambooth-lora-stable-diffusion\"),\n",
+        "    base_model_id=model_id,\n",
+        "    finetuned_lora_model_path=gs_output_dir_dreambooth_lora,\n",
+        "    task=\"text-to-image-lora\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2xT1Bh_XDmZc"
+      },
+      "source": [
+        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed after the above model deployment step succeeds and before you run the next step below. Otherwise you might see a ServiceUnavailable: 503 502:Bad Gateway error when you send requests to the endpoint."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gH23yBtYDmZd"
+      },
+      "outputs": [],
+      "source": [
+        "instances = [\n",
+        "    {\"prompt\": \"A picture of a sks dog on a chair\"},\n",
+        "    {\"prompt\": \"A picture of a sks dog wearing a green hat\"},\n",
+        "    {\"prompt\": \"A picture of a sks dog in front of a tree\"},\n",
+        "    {\"prompt\": \"A picture of a sks dog in a bucket\"},\n",
+        "]\n",
+        "response = endpoint_with_peft_dreambooth.predict(instances=instances)\n",
+        "images = [base64_to_image(image) for image in response.predictions]\n",
+        "image_grid(images)"
+      ]
+    },
+    {
       "cell_type": "markdown",
       "metadata": {
         "id": "ed3795d474b9"
       },
       "source": [
-        "Clean up resources:"
+        "## Clean up"
       ]
     },
     {
@@ -770,11 +1066,16 @@
       },
       "outputs": [],
       "source": [
-        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
-        "\n",
+        "# Undeploy model and delete endpoints.\n",
+        "endpoint_text_to_image.delete(force=True)\n",
+        "endpoint_image_to_image.delete(force=True)\n",
+        "endpoint_with_peft.delete(force=True)\n",
+        "endpoint_with_peft_dreambooth.delete(force=True)\n",
         "# Delete models.\n",
-        "model.delete()"
+        "model_text_to_image.delete()\n",
+        "model_image_to_image.delete()\n",
+        "model_with_peft.delete()\n",
+        "model_with_peft_dreambooth.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_2_1.ipynb

@@ -249,10 +249,10 @@
       "outputs": [],
       "source": [
         "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
       ]
     },
     {