Copy Vertex MG files to notebook folder

2026-09-26 22:51:56 +00:00 · 2024-09-17 23:00:47 +00:00
144 changed files with 15181 additions and 0 deletions
@@ -0,0 +1,253 @@
+# ViT PyTorch vs JAX training benchmarks on Vertex AI Training Platform
+
+Lav Rai, Software Engineer, Google Cloud
+
+Xiang Xu, Software Engineer, Google Cloud
+
+Andreas Steiner, Software Engineer, Google DeepMind
+
+Tao Wang, Software Engineer, Google DeepMind
+
+Alexander Kolesnikov, Research Engineer, Google DeepMind
+
+## Introduction
+
+Many repositories now offer both PyTorch and JAX versions of a model. For
+example, [Hugging Face offers many models such as GPT2, BERT][1]
+etc. Other examples are [OpenLLaMa][2] and [ViT][3]
+models which were first developed in JAX and then their corresponding PyTorch
+versions were made available. **Given both the PyTorch and JAX options for a
+model, it may not be obvious as to which option to choose**. To make such a
+decision, it is important for one to know about the training cost, effectiveness
+and efficiency for each choice.
+
+Apart from the framework choice, the other choice that one faces on Vertex AI
+training platform is the type and count of the accelerators. Although the
+[Vertex AI pricing table][4] lists the price per hour for each
+machine, **one may not know beforehand about the training speed of JAX and
+PyTorch frameworks for different types and count of the accelerators**.
+
+If one has access to some training benchmark numbers for the same model
+under (a) PyTorch and JAX frameworks and (b) for different types and count of
+the accelerators, then it will be easier for them to make a cost effective
+decision. Such a benchmark will also aid the developers in identifying strength
+and weakness of different choices and then figure out recipes to remove those
+weaknesses if possible.
+
+This blog uses the ViT [classification models][5] of varying sizes
+to benchmark the training performance of PyTorch and JAX versions on the Vertex
+AI Platform under different machine configurations. The goal is to:
+
+- Benchmark OSS ViT training for both PyTorch and JAX frameworks.
+- Benchmark OSS ViT L16, H14, g14, and G14 models.
+- Benchmark OSS ViT PyTorch training with A100 GPUs.
+- Benchmark OSS ViT JAX training with A100 GPUs and TPU V3 accelerators.
+
+## Benchmarking setup
+
+This section lays out the benchmarking set up for the [PyTorch][6] and [JAX][7]
+frameworks and provides a reasoning for choosing those settings.
+
+### PyTorch GPU
+
+#### Machine configuration
+
+We run training jobs on [Vertex AI Custom Training][8] using 1
+single node with 8 A100-40GB GPUs.
+
+- Machine type: [a2-highgpu-8g][9]
+- Machine count: 1
+- Accelerator type: [NVIDIA_TESLA_A100 (40GB)][10]
+- Accelerator count: 8
+
+#### Modeling
+
+We benchmark 4 variants of ViT model in different sizes:
+
+- [ViT-L16, 300M params][11]
+- [ViT-H14, 630M params][12]
+- [ViT-g14, 1B params][13]
+- [ViT-G14, 1.8B params][14]
+
+We use the Huggingface [transformers library][15] for ViT L16 and
+H14 variants, and the [TIMM library][16] for ViT g14 and G14
+variants.
+
+#### Dataset
+
+We run training against the [cifar10][17] dataset with 50K training
+images and 10K test images. To factor out network communication overhead for
+data loading, we copy the whole dataset to the local disk then load data from
+the local disk during training.
+
+#### Training parameters
+
+- Trainer
+  - We use [PyTorch Lightning][18] as the trainer for the
+    boilerplate data loading and train loop coding.
+- Precision
+  - Float16
+- Input resolution
+  - 224 x 224
+- Strategy
+  - We use [DDP][19] for models which can be entirely loaded to one
+    GPU, use [Deepspeed-ZeRO][20] otherwise:
+    - ViT-L16: DDP
+    - ViT-H14: DDP
+    - ViT-g14: DDP
+    - ViT-G14: Deepspeed-ZeRO stage-3
+- Batch size
+  - We use the max batch size as power of 2 without CUDA OOM for each model:
+    - ViT-L16: 64 per GPU
+    - ViT-H14: 16 per GPU
+    - ViT-g14: 16 per GPU
+    - ViT-G14: 32 per GPU
+- Compilation
+  - We apply [torch.compile][21] to model whenever it's applicable:
+    - ViT-L16: torch.compile
+    - ViT-H14: torch.compile
+    - ViT-g14: torch.compile
+    - ViT-G14: N/A
+
+### JAX TPU and GPU
+
+#### Machine configuration
+
+All the TPU and GPU training jobs are run on [Vertex AI Custom
+Training][8]. The following machine configurations were used for the
+TPU and GPU experiments:
+
+**Note**: TPU V3 POD requires multi-host supporting training code. For example,
+a 32 core POD runs on 4 hosts with each host using 8 cores.
+
+**Note**: 8 A100 are similar to TPU V3 32 cores in terms of [Vertex AI
+pricing][4].
+
+**Note**: [Each TPU v3 chip has 2 cores which can use 32 GB high-bandwidth
+memory][22] (16 GB per core) so total memory for 32 cores is 16x32 =
+512 GB. Therefore for the same price, TPUs offer more memory than 8 A100-40GB
+GPUs.
+
+#### Modeling
+
+We decided to use an OSS code repository for model implementation. Using an OSS
+repository helps anyone to independently verify the benchmarking results and
+also relate to the results well. For JAX, we selected the
+[Big Vision][23] code repository.
+
+Same as the PyTorch modeling, we benchmark 4 variants of ViT model in different
+sizes:
+
+- [ViT-L16, 300M params][24]
+- [ViT-H14, 630M params][24]
+- [ViT-g14, 1B params][24]
+- [ViT-G14, 1.8B params][24]
+
+**Note**: The [Big Vision code repo][23] has not made the
+checkpoints publicly available for the models larger than the ViT-L16. Therefore
+for the rest of the three variants, the experiments only used random
+initialization for benchmarking the training speed.
+
+#### Dataset
+
+We use training against the [cifar10 TensorFlow dataset][25] with
+50K training images and 10K test images. This dataset is the same as the one
+used for PyTorch experiments except that it is loaded as a TensorFlow dataset.
+Similar to the PyTorch experiments, we copy the whole dataset to the docker
+image to factor out network communication overhead for data loading.
+
+#### Training parameters
+
+- Precision
+  - "bfloat16" setting was used.
+- Input resolution
+  - 224 x 224 after resize (to 448x448) and random crop (to 224x224) before
+    training.
+    - This resolution for training was the same as the PyTorch settings.
+- Strategy
+  - Used DDP for all models except ViT-G14. ViT-G14 used the FSDP strategy.
+- Batch size
+  - We use the max batch size as power of 2 without OOM for each model. The
+    [Benchmarking results][26] section shows the final
+    batch size for each experiment.
+  - Once a maximum batch-size for TPU V3 8 cores was determined, we just scaled
+    it linearly for 32 cores.
+  - Once a maximum batch-size for 1 A100 GPU was determined, we just scaled it
+    linearly for 8 A100 GPUs.
+- Compilation
+  - [jax.jit() compilation][27] is used in JAX codes for efficient
+    execution in XLA.
+- GPU related flags
+  - The following flags are set in the dockerfile for the GPU runs.
+    - Note: _xla_gpu_enable_pipelined_collectives_ is set to false for the
+      ViT-G14 FSDP run.
+
+### Evaluation metric
+
+For both the PyTorch and JAX experiments, the following evaluation metrics are
+collected:
+
+- Throughput: Images-per-second observed for training.
+- Cost: The training-cost-per-epoch (USD).
+
+**Note**: The above metrics are not biased against any framework or machine
+configurations. In addition, these metrics will help one decide the most
+efficient training configurations on Vertex AI.
+
+## Benchmarking results
+
+The lowest cost experiment for each model is marked in **bold** in the last
+column.
+
+![vit_benchmarking_table](images/vit_benchmarking_table.png)
+
+The following bar charts summarize the performance visually:
+
+![vit_training_time](images/vit_training_time.png)
+
+![vit_training_cost](images/vit_training_cost.png)
+
+The following section provides observations and conclusions for these results.
+
+## Observation and Conclusions
+
+- Training with JAX TPU V3 POD with 32 cores costs 33% less than the PyTorch GPU
+  8 A100-40GBs runs.
+- Training with JAX GPU 8 A100-40GBs costs 23% less than the PyTorch GPU 8
+  A100-40GBs runs.
+- JAX TPU V3 POD with 32 cores was 4x faster and slightly more cost-effective
+  than the JAX TPU V3 8 core run for the ViT-large model. This indicates that it
+  might be better to use more cores. The JAX TPU V3 speed scales very well with
+  the number of cores.
+- Cloud TPU VM training speed numbers were the same as the Vertex AI for
+  TPU V3 8 cores. The dataset was copied to the docker in both the cases.
+- The training-cost-per-epoch increases with the model size irrespective of the
+  framework.
+
+[1]: https://github.com/huggingface/transformers/blob/main/examples/research_projects/jax-projects/README.md#quickstart-flax-and-jax-in-transformers
+[2]: https://github.com/openlm-research/open_llama
+[3]: https://github.com/google-research/vision_transformer
+[4]: https://cloud.google.com/vertex-ai/pricing#custom-trained_models
+[5]: https://arxiv.org/abs/2010.11929
+[6]: #pytorch-gpu
+[7]: #jax-tpu-and-gpu
+[8]: https://cloud.google.com/vertex-ai/docs/training/overview
+[9]: https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types
+[10]: https://cloud.google.com/vertex-ai/docs/training/configure-compute#specifying_gpus
+[11]: https://huggingface.co/google/vit-large-patch16-224-in21k
+[12]: https://huggingface.co/google/vit-huge-patch14-224-in21k
+[13]: https://github.com/huggingface/pytorch-image-models/blob/v0.9.2/timm/models/vision_transformer.py#L1308
+[14]: https://github.com/huggingface/pytorch-image-models/blob/v0.9.2/timm/models/vision_transformer.py#L1312
+[15]: https://huggingface.co/docs/transformers/main/model_doc/vit#transformers.ViTModel
+[16]: https://github.com/huggingface/pytorch-image-models
+[17]: https://huggingface.co/datasets/cifar10
+[18]: https://lightning.ai/docs/pytorch/stable/
+[19]: https://pytorch.org/docs/stable/notes/ddp.html
+[20]: https://www.deepspeed.ai/tutorials/zero/
+[21]: https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html
+[22]: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v3
+[23]: https://github.com/google-research/big_vision
+[24]: https://screenshot.googleplex.com/BximJgxsgvBVu38
+[25]: https://www.tensorflow.org/datasets/catalog/cifar10
+[26]: #benchmarking-results
+[27]: https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html
@@ -0,0 +1,188 @@
+# Benchmark report on hyperparameter tuning the OpenLLaMA models on Google Cloud Vertex Model Garden
+
+Changyu Zhu, Software Engineer, Google Cloud
+
+Dustin Luong, Software Engineer, Google Cloud
+
+Gary Wei, Software Engineer, Google Cloud
+
+Genquan Duan, Software Engineer, Google Cloud
+
+## Introduction
+
+Fine-tuning of LLMs can be non-trivial to find an optimal configuration of
+machine types, training parameters, and other hyperparameters that achieves a
+good balance between cost efficiency and model performance. To facilitate users
+in conducting tuning experiments, this report benchmarks fine-tuning OpenLLaMA
+models with [Vertex AI Hyperparameter Tuning Service](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview), demonstrating both efficiency
+and effectiveness. Similar hyperparameter tuning techniques can apply to other models as well.
+
+## Key takeaways
+
+- **The hyperparameter tuning service finds good parameters**: The best model found by the hyperparameter tuning service has an average improvement of around 4% in accuracy in *ARC*, *HellaSwag*, and *TruthfulQA* datasets, while only tuning the learning rate.
+
+- **Hyperparameter tuning works with QLoRA on limited resources**: 4bit QLoRA is sufficient for hyperparameter tuning to find a set of good parameters. In this way, all OpenLLaMA models can run on 1 single `NVIDIA_L4` GPU. It is also possible to train for more steps on the good parameters discovered by hyperparameter tuning, avoiding the waste of computing resources on fine-tuning with suboptimal hyperparameters.
+
+- **Hyperparameter tuning is cost-effective**: While `NVIDIA_L4` is slower than `NVIDIA_TESLA_V100`, it costs less and avoids the overhead of multi-GPU training since it has more GPU memory. Finding a good 3B/7B/13B OpenLLaMA model costs $28.5671, $47.8016, and $87.9208, respectively.
+
+## Benchmarking setup
+
+This section describes the experiment setup of the hyperparameter tuning experiments. The default tuning parameters are:
+
+### Machine configuration
+
+- Machine type: g2-standard-8
+- Machine count: 1
+- Accelerator type: NVIDIA_L4
+- Accelerator count: 1
+
+### Modeling
+
+We benchmark all 3 OpenLLaMA models:
+
+- [open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b)
+- [open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b)
+- [open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b)
+
+We use the Huggingface [PEFT](https://github.com/huggingface/peft) library for fine-tuning.
+
+### Training dataset
+
+We use the dataset [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) loaded directly via Huggingface.
+
+### Training parameters
+
+The set of training parameters used during benchmarking:
+
+- Batch size: 4
+- Precision mode: 4bit QLoRA
+- LoRA rank: 32
+- LoRA alpha: 64
+- Max sequence length: 512
+- Max train steps: 1000
+
+### Evaluation dataset
+
+We use the [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) library injected into the training loop for evaluation. The hyperparameter tuning job will pick the model according to the evaluation metrics.
+
+- Eval task: [ARC Challenge](https://huggingface.co/datasets/ai2_arc)
+- Eval metric: acc_norm
+- Max eval examples: 10000
+
+### Standalone evaluation dataset
+
+After finding the best model with Vertex hyperparameter tuning service, we run standalone evaluations with the model on the following datasets:
+
+- [ARC Challenge](https://huggingface.co/datasets/ai2_arc)
+- [HellaSwag](https://huggingface.co/datasets/Rowan/hellaswag)
+- [TruthfulQA](https://huggingface.co/datasets/EleutherAI/truthful_qa_mc)
+
+### Hyperparameter tuning
+
+We only tune the learning rate hyperparameter. It is considered a floating point value in the continuous range [1e-5, 1e-4]. We run 8 trials in total, with a parallelism of 1 or 2.
+
+### Code example
+
+The following code example launches an example hyperparameter tuning job of OpenLLaMA 7B model.
+
+```py
+from google.cloud import aiplatform
+from google.cloud.aiplatform import hyperparameter_tuning as hpt
+
+
+TRAIN_DOCKER_URI = 'us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231130_0936_RC00'
+output_dir = "gs://path/to/output/dir"
+base_model_id = "openlm-research/open_llama_7b"
+dataset_name = "timdettmers/openassistant-guanaco"
+hpt_precision_mode = "4bit"
+machine_type = "g2-standard-8"
+accelerator_type = "NVIDIA_L4"
+accelerator_count = 1
+eval_task = "arc_challenge"
+eval_metric_name = "acc_norm"
+max_steps = 1000
+eval_limit = 10000
+
+flags = {
+    "learning_rate": 1e-5,
+    "precision_mode": hpt_precision_mode,
+    "task": "instruct-lora",
+    "pretrained_model_id": base_model_id,
+    "output_dir": output_dir,
+    "warmup_steps": 10,
+    "max_steps": max_steps,
+    "lora_rank": 32,
+    "lora_alpha": 64,
+    "lora_dropout": 0.05,
+    "dataset_name": dataset_name,
+    "eval_steps": max_steps + 1,  # Only evaluates at the end.
+    "eval_tasks": eval_task,
+    "eval_limit": eval_limit,
+    "eval_metric_name": eval_metric_name,
+}
+worker_pool_specs = [
+    {
+        "machine_spec": {
+            "machine_type": machine_type,
+            "accelerator_type": accelerator_type,
+            "accelerator_count": accelerator_count,
+        },
+        "replica_count": 1,
+        "container_spec": {
+            "image_uri": TRAIN_DOCKER_URI,
+            "args": ["--{}={}".format(k, v) for k, v in flags.items()],
+        },
+    }
+]
+metric_spec = {"model_performance": "maximize"}
+parameter_spec = {
+    "learning_rate": hpt.DoubleParameterSpec(
+        min=1e-5, max=1e-4, scale="linear"
+    ),
+}
+
+train_job = aiplatform.CustomJob(
+    display_name=job_name,
+    worker_pool_specs=worker_pool_specs,
+    staging_bucket=STAGING_BUCKET,
+)
+
+train_hpt_job = aiplatform.HyperparameterTuningJob(
+    display_name=f"{job_name}_hpt",
+    custom_job=train_job,
+    metric_spec=metric_spec,
+    parameter_spec=parameter_spec,
+    max_trial_count=8,
+    parallel_trial_count=2,
+)
+
+train_hpt_job.run()
+```
+
+## Benchmark results
+
+### Fine-tuning cost
+
+The fine-tuning cost is calculated from `us-central1` pricing and may be subject to changes.
+
+| Model         | Train time | Trials | Parallel Trials | Hourly cost | Cost     | Eval acc_norm (ARC-Challenge) |
+|---------------|------------|--------|-----------------|-------------|----------|-------------------------------|
+| OpenLLaMA 3B  | 16 hrs     | 8      | 2               | $1.7072     | $28.5671 | 39.9%                         |
+| OpenLLaMA 7B  | 28 hrs     | 8      | 2               | $1.7072     | $47.8016 | 45.8%                         |
+| OpenLLaMA 13B | 103 hrs    | 8      | 1               | $0.8536     | $87.9208 | 47.6%                         |
+
+### Fine-tuning performance
+
+Here are the evaluation results of the best model found by hyperparameter tuning, compared with the baseline model. The column `Eval acc_norm` is calculated during training, which is always lower than that during standalone evaluation, because the model is loaded and evaluated at a lower precision (4bit during training / float16 during standalone evaluation).
+
+| Model         | Eval acc_norm (ARC-Challenge) | ARC    | hellaswag | Truthfulqa_mc | ∆ARC   | ∆Hellaswag | ∆Truthfulqa_mc | ∆Average |
+|---------------|-------------------------------|--------|-----------|---------------|--------|------------|----------------|----------|
+| OpenLLaMA 3B  | 39.9%                         | 41.47% | 69.97%    | 38.31%        | +1.62% | +7.32%     | +3.34%         | +4.09%   |
+| OpenLLaMA 7B  | 45.8%                         | 49.83% | 75.53%    | 41.53%        | +2.82% | +3.55%     | +6.68%         | +4.35%   |
+| OpenLLaMA 13B | 47.6%                         | 52.20% | 78.90%    | 44.27%        | +1.01% | +3.67%     | +6.19%         | +3.62%   |
+
+## Related documents
+
+1. [Benchmark report on fine tuning the OpenLLaMA 7B model on Google Cloud Vertex Model Garden
+](
+https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/community-content/vertex_model_garden/benchmarking_reports/pytorch_openllama_7b_finetune_benchmark_report.md)
@@ -0,0 +1,218 @@
+# Benchmark Stable Diffusion v1-5 Fine Tuning and Serving With Google Cloud Vertex Model Garden
+
+Dustin Luong, Software Engineer, Google Cloud
+Gary Wei, Software Engineer, Google Cloud
+Changyu Zhu, Software Engineer, Google Cloud
+Genquan Duan, Software Engineer, Google Cloud
+
+## Introduction
+[The public notebook][1] shows the full examples of fine tuning and serving of Stable diffusion v1-5. [The github repo][2] contains examples of building training and serving dockers for Google Cloud Vertex Model Garden. This report benchmarks Stable diffusion v1-5 fine tuning and serving in Google Cloud Vertex AI, showing both efficiencies and effectiveness.
+
+### Benchmark Highlights
+- Fine tuning
+  - Stable diffusion v1-5 with LoRA and Gradient checkpointing only requires ~10G GPU memory. Larger batch sizes, or larger resolutions require more GPU memories, but not does not change much for different LoRA ranks.
+  - The fine tuning speed is fast in ~11 minutes for 1k steps, and costs less than $1 in 1 A100. The fine tuning speed increases with batch sizes, decreases with resolution, but is not affected much by LoRA ranks.
+  - LoRA tunes a few percent (only 0.1% with LoRA rank=8) of all parameters, and the tuned models are very small (only 3.1MB with LoRA rank=8).
+  - Dreambooth+LoRA and Dreambooth can achieve similar performances, but Dreambooth LoRA can require much less GPU.
+  - Increasing batch size, reducing training steps, and increasing learning rate can result in models with the same performance for less cost.
+- Inference
+  - The optimized serving docker pytorch-peft-serve can speed up inference by 2x than current pytorch-diffuser-serve, and support both base models and fine tuned lora models.
+  - The optimized serving docker pytorch-peft-serve can generate 4 512*512 images in 4.1 seconds on 1 V100 and 1.7 seconds on 1 A100.
+
+Benchmark details are below.
+
+## Fine Tuning Benchmarks
+
+### Experiment Setup
+We mainly compare two tuning algorithms:
+- parameter efficient finetuning based on [dreambooth][3] and [LoRA][4] (shorten as Dreambooth+LoRA below)
+- full parameter fine tuning based on [dreambooth][3] (shorten as Dreambooth below)
+
+And then  report benchmark results on GPU memories, tuning parameters, tuning speeds, costs and accuracy, using the public oxford flowers dataset: [train][5] and [test][6], where the column blip_caption as texts, and column image as images. We also benchmark subject and prompt fidelity using the [dataset][7] from the Dreambooth paper.
+
+The default tuning parameters during benchmark are:
+- Hardware: 1 A100 40G
+- batch size: 4
+- lora_rank: 8
+- resolution: 512
+- max_train_steps: 10
+- use_lora: False
+- gradient_checkpointing: False
+
+```
+# Examples to start finetuning dockers.
+MODEL_NAME="runwayml/stable-diffusion-v1-5"
+OUTPUT_DIR=<OUTPUT_DIR>
+INSTANCE_DATA_DIR=<INSTANCE_DATA_DIR>
+INSTANCE_PROMPT=<INSTANCE_PROMPT>
+IMAGE="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train"
+docker run \
+  --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=0 \
+  --rm --name "test_gpu" \
+  -it ${IMAGE} \
+  --task=text-to-image-dreambooth-lora-peft \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --resolution=512 \
+  --instance_data_dir=$INSTANCE_DATA_DIR \
+  --instance_prompt=$INSTANCE_PROMPT \
+  --train_batch_size=4 \
+  --max_train_steps=10 \
+  --output_dir=${OUTPUT_DIR} \
+  --use_lora \
+  --lora_r=8 \
+  --gradient_checkpointing
+```
+
+### GPU Memories
+Many various factors will impact GPU memory usages. In this benchmark, we mainly benchmark with different finetuning algorithms, batch sizes, lora rank, resolution, and then recommended max batch size on different GPUs.
+
+
+![sd_v1-5_peak_gpu_algorithm](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_algorithm.png)
+
+![sd_v1-5_peak_gpu_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_batch_size.png)
+
+![sd_v1-5_peak_gpu_lora_rank](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_lora_rank.png)
+
+![sd_v1-5_peak_gpu_resolution](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_resolution.png)
+
+- LoRA tuning reduced about 47% peak RAM and 42% peak VRAM for GPU memory, compared to full parameter fine tuning.
+- Gradient checkpointing decreases about 1% peak RAM and 31% peak VRAM for GPU memory further, compared without gradient checkpointing.
+- The GPU memory does not change much for different LoRA ranks.
+- Larger batch sizes require more GPU memories.
+- Larger resolutions require more GPU memories.
+- Dreambooth+LoRA+Gradient_Checkpointing can support max batch size as 32, or max resolution as 2048, but Dreambooth can only support max batch size as 8, or max resolution as 1024.
+
+### Fine Tuning Parameters
+This section shows the percentage of trainable parameters, and tuned model sizes.
+
+- LoRA tunes quite a few percent (only 0.1% with LoRA rank=8) of all parameters, and the tuned models are very small (only 3.1MB with LoRA rank=8).
+
+| LoRA Rank | Trainable parameters | Total parameters | Trainable Parameter Percentage | Fine tuned model size (MB) |
+|---|---|---|---|---|
+| 4 | 398592 | 859919556 | 0.05% | 1.57 |
+|8 | 797184 | 860318148 | 0.09% | 3.09 |
+| 16 | 1594368| 861115332| 0.19%| 6.13|
+| 32| 3188736| 862709700| 0.37%| 12.21|
+### Fine Tuning Speed And Costs
+Fine tuning speeds and costs are affected by many different factors, such as batch size, tuning parameters, image resolutions, GPUs, and datasets. In order to make the report easy to understand, we set the following values in this section:
+- Hardware: 1 A100 40G
+- use_lora: True
+- gradient_checkpointing: True
+
+![sd_v1-5_training_speed_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_speed_batch_size.png)
+
+![sd_v1-5_training_speed_lora_rank](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_speed_lora_rank.png)
+
+![sd_v1-5_training_speed_resolution](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_speed_resolution.png)
+
+![sd_v1-5_training_cost_max_steps](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_cost_max_steps.png)
+
+- The fine tuning speed increases with batch sizes, decreases with resolution, but is not affected much by LoRA ranks.
+- The fine tuning speed is about 11 minutes for 1k steps, and costs less than $1 in 1 A100.
+
+### Fine Tuning Quality
+In this benchmark, we mainly benchmark Dreambooth and Dreambooth+LoRA to compare fine tuning quality. We compare [subject fidelity scored (DINO)][8], how well the subject is represented in the generated images, and [prompt fidelity scores (CoCa)][9], how well the generated images match the given prompt, for a single subject, a [dog][10] from the dataset released with the original Dreambooth paper. In practice, we recommend saving checkpoints periodically and inspecting validation prompts visually. We fine tuned the unet without fine tuning the text encoder and used the following hyperparameters:
+
+Dreambooth
+- Learning rate: 5e-6
+- Batch size: 1
+
+Dreambooth+LoRA
+- Learning rate: 1e-4
+- Batch size: 1
+
+![sd_v1-5_finetuning_quality_subject_fidelity](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_finetuning_quality_subject_fidelity.png)
+
+![sd_v1-5_finetuning_quality_prompt_fidelity](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_finetuning_quality_prompt_fidelity.png)
+
+- Fine tuning with Dreambooth or Dreambooth+LoRA can result in models with comparable performance. The base model produced images of the class rather than the instance.
+- Dreambooth+LoRA is able to achieve the same subject fidelity score as Dreambooth if trained for more epochs.
+- Increasing the number of training steps results in better subject fidelity but at the cost of prompt fidelity.
+
+### Suggested Max Batch Sizes By Resolutions
+We benchmarked and suggested max batch sizes by resolutions on 1 A100 and 1 V100 as below. This is with LoRA and gradient checkpointing enabled.
+
+![sd_v1-5_batch_size_by_resolution](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_batch_size_by_resolution.png)
+
+### Fine Tuning Cost Optimization
+Increasing batch size allows for more images to be considered at each training step for fine tuning. This allows models to be trained in fewer training steps. In this benchmark, we aim to show how batch size can be increased to reduce training costs while still preserving subject and prompt fidelity.
+
+Since the training dataset consists of 5 images, we train with a batch size of 5 and reduce the number of training steps from 400 to 80. Doing so results in a model that has not learned the subject since we’ve decreased the number of training steps. Conceptually, the model is taking a more precise step at each iteration, but it is taking fewer steps. To compensate for this, we increased the learning rate from 5e-6 and observed the best results at 1e-5 for full parameter finetuning.
+
+![sd_v1-5_subject_fidelity_batch_size_5_learning_rate](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_subject_fidelity_batch_size_5_learning_rate.png)
+
+![sd_v1-5_prompt_fidelity_batch_size_5_learning_rate](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_prompt_fidelity_batch_size_5_learning_rate.png)
+
+Comparing cost of training the “best” model for batch size 1 vs. batch size 5
+
+![sd_v1-5_cost_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_cost_batch_size.png)
+
+
+| Train method| Training parameters| Sample image| CoCa (prompt fidelity)| DINO (subject fidelity) | Cost of training on A100 |
+|---|---|---|---|---|---|
+| dreambooth| dreambooth, num_train_steps=400, batch_size=1, lr=5e-6| ![dog1](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog1.png) | 0.12215| 0.76531| $0.26 |
+| dreambooth | dreambooth, num_train_steps=80, batch_size=5,lr=1e-5| ![dog2](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog2.png)| 0.12644| 0.74697 | $0.15 |
+| dreambooth-lora| num_train_steps=500, batch_size=1, lr=1e-4, gc|![dog3](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog3.png)| 0.12856| 0.78148 | $0.26|
+| dreambooth-lora | num_train_steps=50, batch_size=5, lr=1e-3, gc | ![dog4](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog4.png) | 0.12566 | 0.75479 | $0.09 |
+
+
+
+A followup question is that since finetuning can be run on a single GPU, should finetuning be run on 1 V100 or A100?
+
+Setup:
+- num_train_steps=800 / batch_size
+- Resolution=512
+
+![sd_v1-5_cost_training_method_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_cost_training_method_batch_size.png)
+- Although V100 has a lower $/hr cost than an A100, the same training setup takes longer. Even given the longer training time, the cost on V100 is still lower.
+- Dreambooth+LoRA enables training with larger batch sizes, however, larger batch sizes will not necessarily mean faster training time.
+- It is possible to fine tune with 1 V100 on 512 resolution with Dreambooth+LoRA.
+- Dreambooth fine tuning must be run on 1 A100 at 512 resolution.
+
+## Inference Benchmarks
+We provide two serving dockers in vertex model garden for stable diffusion:
+- pytorch-diffuser-serve:
+  - us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve
+  - This serving docker only serves base stable diffusion models and does not contain any optimizations yet.
+- pytorch-peft-serve:
+  - us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve
+  - This serving docker can serve base stable diffusion models, and base stable diffusion models with fine tuned lora models, and contains optimization for serving.
+
+We run the two serving dockers on T4/V100/A100 to generate 4 512*512 images, and compare the inference speed without network considerations as:
+
+![sd_v1-5_inference_speed_gpu](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_inference_speed_gpu.png)
+The speed up of optimized pytorch-peft-serve is about 2x than current pytorch-diffuser-serve.
+
+### Serving cost comparison
+
+Pytorch-diffuser-serve (without any optimizations)
+
+| GPU type|  Time required to generate 4 512x512 images | Machine unit price ($ / hour) | Cost per image ($) |
+|---|---|---|---|
+| T4 | 28.6 | 0.4025| 0.00080 |
+| V100 | 8.8 | 2.852| 0.00174|
+| A100 | 4.2 | 4.2245 | 0.00123 |
+
+Pytorch-peft-serve (with optimizations)
+
+| GPU type | Time required to generate 4 512x512 images | Machine unit price ($ / hour) | Cost per image ($) |
+|--- |---|---|---|
+| T4 | 12.6 | 0.4025 | 0.00035 |
+| V100 | 4.1 | 2.852 | 0.00081 |
+|  A100 | 1.7 | 4.2245 | 0.00050 |
+
+- The optimized pytorch-peft-serve has approximately half the price per image, compared with the un-optimized pytorch-diffuser-serve.
+- Serving the model with a T4 is most cost effective, however, serving with an A100 still has the best throughput and fastest predictions.
+
+
+
+[1]: https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb
+[2]: https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/community-content/vertex_model_garden/model_oss
+[3]: https://arxiv.org/abs/2208.12242
+[4]: https://arxiv.org/abs/2106.09685
+[5]: https://huggingface.co/datasets/Multimodal-Fatima/OxfordFlowers_train
+[6]: https://huggingface.co/datasets/Multimodal-Fatima/OxfordFlowers_test_facebook_opt_6.7b_Attributes_ns_6149
+[7]: https://github.com/google/dreambooth
+[8]: https://arxiv.org/abs/2104.14294
+[9]: https://arxiv.org/abs/2205.01917
+[10]: https://github.com/google/dreambooth/tree/main/dataset/dog6
@@ -0,0 +1,50 @@
+# Dockerfile for serving dockers with AutoGluon.
+#
+# To build:
+# docker build -f model_oss/autogluon/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/pytorch:2.1.2-cuda11.8-cudnn8-runtime
+
+USER root
+
+# AutoGluon might require libgomp for some dependencies.
+RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        libgomp1
+
+# Install AutoGluon and other dependencies.
+RUN pip install --upgrade pip
+RUN pip install autogluon==1.0.0
+RUN pip install flask==3.0.0
+
+# Dependencies needed to work with GCS.
+RUN pip install absl-py==2.0.0
+RUN pip install google-cloud-storage==2.7.0
+
+# Copy scripts into the container.
+COPY model_oss/autogluon /autogluon
+COPY model_oss/util /autogluon/util
+WORKDIR /autogluon
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+RUN wget https://github.com/pallets/flask/blob/main/LICENSE.rst
+
+# Expose the port the app runs on.
+EXPOSE 8501
+
+# Set the working directory to a specific path for consistency.
+WORKDIR /autogluon
+
+# Change to a non-root user for security purposes.
+RUN useradd -m autogluonuser
+USER autogluonuser
+
+# Run Flask application.
+CMD ["python", "serve.py"]
@@ -0,0 +1,36 @@
+# Dockerfile for training dockers with Autogluon.
+#
+# To build:
+# docker build -f model_oss/autogluon/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.1.2-cuda11.8-cudnn8-runtime
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PIP_ROOT_USER_ACTION=ignore
+
+RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+apt-utils \
+curl \
+wget \
+git \
+jq \
+gnupg \
+build-essential \
+tesseract-ocr \
+vim
+
+# Install libraries.
+RUN pip install autogluon==1.0.0
+
+COPY model_oss/autogluon /autogluon
+WORKDIR /autogluon
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENTRYPOINT ["python", "train.py"]
@@ -0,0 +1,87 @@
+r"""AutoGluon serving binary.
+
+This module sets up a Flask web server for serving predictions from a
+trained AutoGluon model. The server exposes two endpoints:
+
+1.  `/ping`: A health check endpoint that returns "pong" to
+    indicate that the server is running.
+2.  `/predict`: An endpoint that accepts POST requests with JSON content.
+    Each request should contain one or more instances for which the
+    predictions are desired. The endpoint returns the predictions and
+    associated probabilities in a JSON response.
+
+The server expects an environment variable `model_path` that points to
+the directory where the AutoGluon model artifacts are
+stored. If `model_path` is not provided, it defaults to '/autogluon/models'.
+"""
+
+import json
+import logging
+import os
+
+from autogluon.tabular import TabularPredictor
+import flask
+import pandas as pd
+
+from util import constants
+from util import fileutils
+
+_SUCCESS_STATUS = 200
+_ERROR_STATUS = 500
+_PORT = 8501
+
+app = flask.Flask(__name__)
+# Check the environment variables.
+model_dir = os.getenv('model_path', '/autogluon/models')
+logging.info('Model directory passed by the user is: %s', model_dir)
+# If the model is on GCS then copy it to a local folder first.
+if model_dir.startswith(constants.GCS_URI_PREFIX):
+  gcs_path = model_dir[len(constants.GCS_URI_PREFIX) :]
+  local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
+  logging.info('Download %s to %s', model_dir, local_model_dir)
+  fileutils.download_gcs_dir_to_local(model_dir, local_model_dir)
+  model_dir = local_model_dir
+  logging.info('Local model directory is: %s', model_dir)
+
+
+# Load the predictor at startup.
+predictor = TabularPredictor.load(model_dir)
+
+
+@app.route('/ping', methods=['GET'])
+def ping() -> flask.Response:
+  """Health check route."""
+  return flask.Response('pong', status=_SUCCESS_STATUS)
+
+
+@app.route('/predict', methods=['POST'])
+def predict() -> flask.Response:
+  """Prediction route."""
+  try:
+    # Extract JSON content from the POST request.
+    data = flask.request.get_json(force=True)
+    instances = data.get('instances', [])
+
+    # Convert instances to DataFrame.
+    df_to_predict = pd.DataFrame(instances)
+
+    # Perform prediction.
+    predictions = predictor.predict(df_to_predict).tolist()
+    response = {'predictions': predictions}
+
+    return flask.Response(
+        json.dumps(response),
+        status=_SUCCESS_STATUS,
+        mimetype='application/json',
+    )
+
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    return flask.Response(
+        json.dumps({'error': str(e)}),
+        status=_ERROR_STATUS,
+        mimetype='application/json',
+    )
+
+
+if __name__ == '__main__':
+  app.run(host='0.0.0.0', port=_PORT)
@@ -0,0 +1,144 @@
+"""AutoGluon training binary. """
+
+import argparse
+import json
+from typing import Any
+
+from autogluon.tabular import TabularPredictor
+import pandas as pd
+
+
+class BaseConfig:
+
+  def to_dict(self) -> dict[str, Any]:
+    return {
+        key: value for key, value in self.__dict__.items() if value is not None
+    }
+
+
+class DataConfig(BaseConfig):
+
+  def __init__(self, train_data_path: Any) -> None:
+    self.train_data_path = train_data_path
+
+
+class ProblemConfig(BaseConfig):
+
+  def __init__(self, label: Any, problem_type: Any) -> None:
+    self.label = label
+    self.problem_type = problem_type
+
+
+class EvaluationConfig(BaseConfig):
+
+  def __init__(self, eval_metric: Any) -> None:
+    self.eval_metric = eval_metric
+
+
+class TrainingConfig(BaseConfig):
+  """Config for training."""
+
+  def __init__(
+      self,
+      time_limit: Any,
+      presets: Any,
+      hyperparameters: Any,
+      model_save_path: str,
+  ) -> None:
+    self.time_limit = time_limit
+    self.hyperparameters = hyperparameters
+    self.presets = presets
+    self.model_save_path = model_save_path
+
+
+def parse_args() -> (
+    tuple[DataConfig, ProblemConfig, EvaluationConfig, TrainingConfig]
+):
+  """Parse command line arguments."""
+  parser = argparse.ArgumentParser(description="AutoGluon Tabular Predictor")
+  # Add arguments for each config class
+  parser.add_argument(
+      "--train_data_path",
+      type=str,
+      required=True,
+      help="Path to the input data CSV file.",
+  )
+  parser.add_argument(
+      "--label", type=str, required=True, help="Target variable column name."
+  )
+  parser.add_argument(
+      "--problem_type",
+      type=str,
+      choices=["binary", "multiclass", "regression", "quantile"],
+      default=None,
+      help="Problem type.",
+  )
+  parser.add_argument(
+      "--eval_metric", type=str, default=None, help="Evaluation metric to use."
+  )
+  # Add arguments for TrainingConfig if needed
+  parser.add_argument(
+      "--time_limit",
+      type=int,
+      default=None,
+      help="Time limit in seconds for training.",
+  )
+  parser.add_argument(
+      "--presets",
+      type=str,
+      default="medium_quality",
+      help="Presets used for training ",
+  )
+  parser.add_argument(
+      "--hyperparameters",
+      type=json.loads,
+      default=None,
+      help="Hyperparameter dictionary in JSON format.",
+  )
+  parser.add_argument(
+      "--model_save_path",
+      type=str,
+      default=None,
+      help="Path to save the trained model.",
+  )
+  args = parser.parse_args()
+
+  data_config = DataConfig(train_data_path=args.train_data_path)
+  problem_config = ProblemConfig(
+      label=args.label, problem_type=args.problem_type
+  )
+  eval_config = EvaluationConfig(eval_metric=args.eval_metric)
+  training_config = TrainingConfig(
+      time_limit=args.time_limit,
+      presets=args.presets,
+      hyperparameters=args.hyperparameters,
+      model_save_path=args.model_save_path,
+  )
+
+  return data_config, problem_config, eval_config, training_config
+
+
+def main() -> None:
+  data_config, problem_config, eval_config, training_config = parse_args()
+
+  # Load the training data.
+  data = pd.read_csv(data_config.train_data_path)
+
+  # Create a TabularPredictor.
+  predictor = TabularPredictor(
+      label=problem_config.label,
+      eval_metric=eval_config.eval_metric,
+      path=training_config.model_save_path,
+  )
+
+  # Fit the model
+  predictor.fit(
+      data,
+      presets=training_config.presets,
+      time_limit=training_config.time_limit,
+      hyperparameters=training_config.hyperparameters,
+  )
+
+
+if __name__ == "__main__":
+  main()
@@ -0,0 +1,25 @@
+# The provided content is a configuration file for the ZipNeRF
+# PyTorch implementation.
+
+# Sets the name of the experiment to 'test'.
+Config.exp_name = 'test'
+# Specifies the dataset loader, in this case, 'llff' for light field.
+Config.dataset_loader = 'llff'
+# Defines the near and far clipping planes for the camera view.
+Config.near = 0.2
+Config.far = 1e6
+# Image downsampling.
+Config.factor = 4
+
+# For the model configurations.
+Model.raydist_fn = 'power_transformation'
+Model.opaque_background = True
+
+# Disables the computation of density normals and RGB values, and sets
+# the grid level dimension to 1 for PropMLP.
+PropMLP.disable_density_normals = True
+PropMLP.disable_rgb = True
+PropMLP.grid_level_dim = 1
+
+# Disable density normals for NerfMLP
+NerfMLP.disable_density_normals = True
@@ -0,0 +1,21 @@
+# The provided content is a configuration file for Generative
+# Latent Optimization (GLO) vectors in the Pytorch implemnetation of ZipNeRF.
+
+# Specifies the dataset loader, in this case, 'llff' for light field.
+Config.dataset_loader = 'llff'
+# Defines the near and far clipping planes for the camera view.
+Config.near = 0.2
+Config.far = 1e6
+# Image downsampling.
+Config.factor = 4
+
+# For the model configurations.
+Model.raydist_fn = 'power_transformation'
+Model.num_glo_features = 128
+Model.opaque_background = True
+
+PropMLP.disable_density_normals = True
+PropMLP.disable_rgb = True
+PropMLP.grid_level_dim = 1
+
+NerfMLP.disable_density_normals = True
@@ -0,0 +1,18 @@
+# The provided content is a configuration file running ZipNeRF
+# training on 8 gpu machine.
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
@@ -0,0 +1,120 @@
+# Dockerfile for ZipNeRF base image.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_base.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.1.0-cuda11.8-cudnn8-devel
+
+USER root
+
+ARG COLMAP_GIT_COMMIT=main
+ARG CUDA_ARCHITECTURES=60;70;75;80;86
+
+# Prevent stop building ubuntu at time zone selection.
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update -y --allow-releaseinfo-change && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        g++ \
+        wget \
+        vim \
+        bash \
+        cmake \
+        imagemagick \
+        ninja-build \
+        build-essential \
+        libboost-program-options-dev \
+        libboost-filesystem-dev \
+        libboost-graph-dev \
+        libboost-system-dev \
+        libeigen3-dev \
+        libflann-dev \
+        libfreeimage-dev \
+        libmetis-dev \
+        libgoogle-glog-dev \
+        libgtest-dev \
+        libsqlite3-dev \
+        libglew-dev \
+        qtbase5-dev \
+        libqt5opengl5-dev \
+        libcgal-dev \
+        libceres-dev \
+        git \
+        git-lfs \
+        python3-cffi \
+        python3-cryptography \
+        libffi-dev \
+        python-dev
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# 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 deps and install gsutil.
+RUN pip install gsutil==5.27
+
+# When building colmap in colab, the link error "undefined reference.
+# to '_glapi_tls_Current'" happens. A solution is to install "libglvnd"
+# as described in this page https://github.com/colmap/colmap/issues/1271.
+RUN git clone --depth 1 --branch v1.7.0 https://github.com/NVIDIA/libglvnd && \
+    apt-get install -y libxext-dev libx11-dev x11proto-gl-dev && \
+    cd libglvnd/  && \
+    apt-get install -y autoconf automake libtool && \
+    apt-get install -y libffi-dev && \
+    ./autogen.sh && \
+    ./configure && \
+    make  -j4 && \
+    make install
+
+RUN apt remove nvidia-cuda-toolkit -y \
+    nvidia-cuda-toolkit \
+    nvidia-cuda-toolkit-gcc
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+
+ENV CUDA_HOME=/usr/local/cuda
+
+RUN git clone --branch main https://github.com/SuLvXiangXin/zipnerf-pytorch.git
+# Set current directory to the downloaded 'zipnerf-pytorch' repository.
+WORKDIR ./zipnerf-pytorch
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard 4de3d21ebb9e15412d36951b56e2d713fddd812b
+COPY model_oss/cloudnerf/requirements.txt requirements.txt
+RUN pip install -r requirements.txt
+
+# Install gridencoder extensions and nvdiffrast (for textured mesh).
+RUN cd .. && \
+    TORCH_CUDA_ARCH_LIST="6.0 7.0 7.5 8.0 8.6+PTX" CXX=g++ pip install ./zipnerf-pytorch/gridencoder
+
+# Install cuda version of torch_scatter.
+RUN pip install torch-scatter==2.1.2 -f https://data.pyg.org/whl/torch-2.0.1+cu118.html
+RUN pip install google-cloud-aiplatform==1.25.0
+RUN pip install google-cloud-storage==2.9.0
+
+# Build and install COLMAP.
+RUN git clone  --depth 1 --branch 3.8 https://github.com/colmap/colmap.git
+RUN cd colmap && \
+    git fetch https://github.com/colmap/colmap.git ${COLMAP_GIT_COMMIT} && \
+    mkdir build && \
+    cd build && \
+    cmake .. -GNinja -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES} && \
+    ninja && \
+    ninja install && \
+    cd .. && rm -rf colmap
+
+RUN git clone  --depth 1 --branch v1.0.2 https://github.com/dranjan/python-plyfile.git
+
+RUN sed -i "20 i\sys.path.append('/workspace/zipnerf-pytorch/internal/pycolmap')" /workspace/zipnerf-pytorch/internal/datasets.py
+RUN sed -i "21 i\sys.path.append('/workspace/zipnerf-pytorch/internal/pycolmap/pycolmap')" /workspace/zipnerf-pytorch/internal/datasets.py
@@ -0,0 +1,16 @@
+# Dockerfile for ZipNeRF COLMAP image calibration.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/cloudnerf_pytorch_calibrate.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 us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
+
+COPY model_oss/cloudnerf/local_colmap_and_resize.sh /workspace/zipnerf-pytorch/scripts/local_colmap_and_resize.sh
+
+WORKDIR /workspace/zipnerf-pytorch/
+
+ENTRYPOINT ["bash","scripts/local_colmap_and_resize.sh"]
@@ -0,0 +1,22 @@
+# Dockerfile for ZipNeRF rendering.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_render.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 us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
+
+COPY model_oss/cloudnerf/render.sh /workspace/zipnerf-pytorch/scripts/render.sh
+COPY model_oss/cloudnerf/configs/360.gin /workspace/zipnerf-pytorch/configs/360.gin
+COPY model_oss/cloudnerf/configs/360_glo.gin /workspace/zipnerf-pytorch/configs/360_glo.gin
+COPY model_oss/cloudnerf/configs/accelerate_config.yaml /root/.cache/huggingface/accelerate/default_config.yaml
+RUN sed -i '324s/.*/            keyframe_names = fp.read().splitlines()/' /workspace/zipnerf-pytorch/internal/camera_utils.py
+
+ENV PYTHONPATH "${PYTHONPATH}:/workspace/zipnerf-pytorch/util"
+
+WORKDIR /workspace/zipnerf-pytorch/
+
+ENTRYPOINT ["bash", "scripts/render.sh"]
@@ -0,0 +1,21 @@
+# Dockerfile for ZipNeRF training.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_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 us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
+
+COPY model_oss/cloudnerf/train.sh /workspace/zipnerf-pytorch/scripts/train.sh
+COPY model_oss/cloudnerf/configs/360.gin /workspace/zipnerf-pytorch/configs/360.gin
+COPY model_oss/cloudnerf/configs/360_glo.gin /workspace/zipnerf-pytorch/configs/360_glo.gin
+COPY model_oss/cloudnerf/configs/accelerate_config.yaml /root/.cache/huggingface/accelerate/default_config.yaml
+
+ENV PYTHONPATH "${PYTHONPATH}:/workspace/zipnerf-pytorch/util"
+
+WORKDIR /workspace/zipnerf-pytorch/
+
+ENTRYPOINT ["bash", "scripts/train.sh"]
@@ -0,0 +1,144 @@
+#!/bin/bash
+# This script runs colmap for scale invariant feature (SIFT) extraction and
+# matching to map camera extrinsics and intrinsics values for ZipNeRF,
+# given a folder of images and videos
+# from a GCS bucket. It uses ffmepg to extract an image from a video at
+# 1fps. The folder can contain images or videos. If both images and videos
+# are present, the extracted frames from the videos is added to the images
+# to create the final combined image dataset.
+# vv-docker:google3-begin(internal)
+# TODO(b/314042136): Specify cloudnerf colmap fps.
+# vv-docker:google3-end
+
+# Initialize variables.
+use_gpu=1  # Default to 1 (assuming the docker is run on a machine with GPU)
+gcs_dataset_path=""
+gcs_experiment_path=""
+camera=""
+
+# This loop processes command-line arguments for configuring the container.
+# It supports arguments for GPU usage, dataset and experiment paths,
+# and camera type.
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    -use_gpu)
+      use_gpu="$2"
+      if ! [[ $use_gpu =~ ^[0-9]+$ ]]; then
+        echo "Error: -use_gpu must be an integer."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    -gcs_dataset_path)
+      gcs_dataset_path="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -gcs_experiment_path)
+      gcs_experiment_path="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -camera)
+      camera="$2"
+      if [[ $camera != "OPENCV" && $camera != "OPENCV_FISHEYE" ]]; then
+        echo "Error: -camera must be either 'OPENCV' or 'OPENCV_FISHEYE'."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    *)    # unknown option
+      echo "Unknown option: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+local_folder="dataset_content"
+images_folder="dataset_images"
+images_subfolder="images"
+output_folder="$images_folder/$images_subfolder"
+
+# Create the local folder if it doesn't exist
+mkdir -p "$local_folder"
+mkdir -p "$output_folder"
+
+# Download the content from the GCS URI
+gsutil -m cp -r "$gcs_dataset_path"/* "$local_folder/"
+
+# Process files in the local folder
+for file in "$local_folder"/*; do
+  if [[ -f "$file" ]]; then
+    # Check if the file is an image (e.g., jpg, png, etc.)
+    if file --mime-type "$file" | grep -q "image"; then
+      # Copy the image to the "images" subfolder within the "dataset_images" folder
+      cp "$file" "$output_folder/$(basename "$file")"
+    elif file --mime-type "$file" | grep -q "video"; then
+      # Use FFmpeg to extract an image every 30 frames from the video
+      ffmpeg -i "$file" -vf "select='not(mod(n,30))'" "$output_folder/$(basename "$file" ."${file##*.}")_%03d.jpg"
+    else
+      echo "Skipping unsupported file: $file"
+    fi
+  fi
+done
+
+# Run COLMAP Feature extraction
+colmap feature_extractor \
+    --database_path "$local_folder"/database.db \
+    --image_path "$output_folder" \
+    --ImageReader.single_camera 1 \
+    --ImageReader.camera_model "$camera" \
+    --SiftExtraction.use_gpu "$use_gpu"
+
+# Run COLMAP Feature matching
+colmap exhaustive_matcher \
+    --database_path "$local_folder"/database.db \
+    --SiftMatching.use_gpu "$use_gpu"
+
+# Bundle adjustment. The default Mapper tolerance is unnecessarily large,
+# decreasing it speeds up bundle adjustment steps.
+mkdir -p "$local_folder"/sparse
+colmap mapper \
+    --database_path "$local_folder"/database.db \
+    --image_path "$output_folder" \
+    --output_path "$local_folder"/sparse \
+    --Mapper.ba_global_function_tolerance=0.000001
+
+# Downsample images at 1/2, 1/4, 1/8 scales. Save feature matching to
+# sqlite database.
+# All input and output images:
+#   $gcs_dataset_path
+#   $gcs_experiment_path/data/images
+# Downsampled output images:
+#   $gcs_experiment_path/data/images_2/
+#   $gcs_experiment_path/data/images_4/
+#   $gcs_experiment_path/data/images_8/
+# COLMAP sparse reconstruction files: project.ini, images.bin,
+# cameras.bin, points3D.bin
+#   $gcs_experiment_path/data/sparse/0/
+cp -r "$output_folder" "$images_folder"/images_2
+pushd "$images_folder"/images_2
+ls | xargs -P 8 -I {} mogrify -resize 50% {}
+popd
+gsutil -m cp -r "$images_folder"/images_2/*  "$gcs_experiment_path"/data/images_2
+
+cp -r "$output_folder" "$images_folder"/images_4
+pushd "$images_folder"/images_4
+ls | xargs -P 8 -I {} mogrify -resize 25% {}
+popd
+gsutil -m cp -r "$images_folder"/images_4/*  "$gcs_experiment_path"/data/images_4
+
+cp -r "$output_folder" "$images_folder"/images_8
+pushd "$images_folder"/images_8
+ls | xargs -P 8 -I {} mogrify -resize 12.5% {}
+popd
+gsutil -m cp "$images_folder"/images_8/*  "$gcs_experiment_path"/data/images_8
+
+# Copy images and sparse reconstruction files to gcs experiment folder.
+gsutil -m cp "$images_folder"/images/*  "$gcs_experiment_path"/data/images
+gsutil -m cp -r "$local_folder"/sparse  "$gcs_experiment_path"/data
+gsutil -m cp "$local_folder"/database.db  "$gcs_experiment_path"/data
+
+echo "Processing complete."
@@ -0,0 +1,117 @@
+#!/bin/bash
+# This script runs rendering for ZipNeRF given an experiment folder
+# from a GCS bucket with colmap dataset.
+
+# Initialize associative array for arguments.
+declare -A args
+
+# vv-docker:google3-begin(internal)
+# TODO(b/311468174): Pass gin config file from gcs bucket.
+# vv-docker:google3-end
+# Function to parse named arguments.
+parse_args() {
+  while [[ $# -gt 0 ]]; do
+    key="$1"
+    case $key in
+      -gcs_experiment_path|-gin_config_file|-gcs_keyframes_file)
+        args[$key]="$2"
+        shift # past argument
+        shift # past value
+        ;;
+      -training_job_name)
+        training_job_name="$2"
+        shift # past argument
+        shift # past value
+        ;;
+      -rendering_job_name)
+        rendering_job_name="$2"
+        shift # past argument
+        shift # past value
+        ;;
+      -render_path_frames|-factor|-render_video_fps)
+        args[$key]="$2"
+        if ! [[ ${args[$key]} =~ ^[0-9]+$ ]]; then
+          echo "Error: $key must be an integer."
+          exit 1
+        fi
+        shift # past argument
+        shift # past value
+        ;;
+      *)
+        echo "Unknown option: $1" >&2
+        exit 1
+        ;;
+    esac
+  done
+}
+
+# Function to create a directory if it doesn't exist.
+create_dir_if_not_exists() {
+  local dir_path=$1
+  if [[ ! -d "$dir_path" ]]; then
+    echo "Creating folder: $dir_path"
+    mkdir "$dir_path"
+  else
+    echo "Folder $dir_path already exists."
+  fi
+}
+
+# Function to launch rendering.
+launch_rendering() {
+  local keyframes_file=$1
+  local render_bindings=(
+    "--gin_configs=${args[-gin_config_file]}"
+    "--gin_bindings=Config.data_dir='${DATASET_PATH}'"
+    "--gin_bindings=Config.exp_name='${EXPERIMENT}'"
+    "--gin_bindings=Config.render_path=True"
+    "--gin_bindings=Config.render_path_frames=${args[-render_path_frames]}"
+    "--gin_bindings=Config.render_video_fps=${args[-render_video_fps]}"
+    "--gin_bindings=Config.factor=${args[-factor]}"
+  )
+
+  if [[ -n $keyframes_file ]]; then
+    render_bindings+=("--gin_bindings=Config.render_spline_keyframes='${keyframes_file}'")
+  fi
+
+  accelerate launch render.py "${render_bindings[@]}"
+}
+
+# Parse arguments.
+parse_args "$@"
+
+# Extract folder names and paths.
+scene_folder_name=$(basename "${args[-gcs_experiment_path]}")
+local_dataset_path="local_dataset"
+local_experiment_path="exp"
+exp_folder_name=$(basename "${args[-gcs_experiment_path]}")
+DATASET_PATH="$local_experiment_path/$exp_folder_name/data"
+CHECKPOINTS_PATH="$local_experiment_path/$exp_folder_name/checkpoints"
+OUTPUT_RENDER_PATH="$local_experiment_path/$scene_folder_name/render"
+EXPERIMENT=$exp_folder_name
+
+# Create necessary directories.
+create_dir_if_not_exists "$local_dataset_path"
+create_dir_if_not_exists "$local_experiment_path"
+create_dir_if_not_exists "$local_experiment_path/$exp_folder_name"
+create_dir_if_not_exists "$CHECKPOINTS_PATH"
+
+# Create the file log_render.txt in the exp folder.
+touch "$local_experiment_path/$exp_folder_name/log_render.txt"
+
+# Copy experiment from GCS bucket to local
+gsutil -m cp -r "${args[-gcs_experiment_path]}/data" "$local_experiment_path/$exp_folder_name" || exit 1
+gsutil -m cp -r "${args[-gcs_experiment_path]}/checkpoints/${training_job_name}/*" "$CHECKPOINTS_PATH" || exit 1
+
+# Check and copy keyframes file.
+if [[ -n ${args[-gcs_keyframes_file]} ]]; then
+  keyframes_file_basename=$(basename "${args[-gcs_keyframes_file]}")
+  local_keyframes_file="$local_dataset_path/$keyframes_file_basename"
+  gsutil cp "${args[-gcs_keyframes_file]}" "$local_keyframes_file" || exit 1
+  echo "Local keyframe file: $local_keyframes_file"
+  launch_rendering "$local_keyframes_file"
+else
+  launch_rendering ""
+fi
+
+# Copy rendered data back to GCS.
+gsutil -m cp -r "$OUTPUT_RENDER_PATH" "${args[-gcs_experiment_path]}/render/${rendering_job_name}"
@@ -0,0 +1,24 @@
+--find-links https://download.pytorch.org/whl/torch_stable.html
+
+torch==2.2.0
+numpy==1.26.1
+absl_py==2.0.0
+accelerate==0.24.0
+gin_config==0.5.0
+imageio==2.31.6
+imageio-ffmpeg==0.4.9
+matplotlib==3.8.0
+mediapy==1.1.9
+ninja==1.11.1.1
+opencv_contrib_python==4.8.1.78
+opencv_python==4.8.1.78
+Pillow==10.3.0
+rawpy==0.18.1
+scipy==1.11.3
+scikit-image==0.22.0
+scikit-learn==1.5.0
+tensorboard==2.15.0
+tensorboardX==2.6.2.2
+tqdm==4.66.3
+trimesh==4.0.1
+xatlas==0.0.8
@@ -0,0 +1,94 @@
+#!/bin/bash
+
+# Initialize variables.
+training_job_name=""
+gcs_experiment_path=""
+gin_config_file="configs/360.gin"
+factor=4
+max_training_steps=25000
+
+# Parse named arguments.
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    -training_job_name)
+      training_job_name="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -gcs_experiment_path)
+      gcs_experiment_path="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -gin_config_file)
+      gin_config_file="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -factor)
+      factor="$2"
+      if ! [[ $factor =~ ^[0-9]+$ ]]; then
+        echo "Error: -factor must be an integer."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    -max_training_steps)
+      max_training_steps="$2"
+      if ! [[ $max_training_steps =~ ^[0-9]+$ ]]; then
+        echo "Error: -max_training_steps must be an integer."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    *)    # unknown option
+      echo "Unknown option: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+# Function to create a directory if it doesn't exist.
+create_dir_if_not_exists() {
+  local dir_path=$1
+  if [[ ! -d "$dir_path" ]]; then
+    echo "Creating folder: $dir_path"
+    mkdir "$dir_path"
+  else
+    echo "Folder $dir_path already exists."
+  fi
+}
+
+# Extract folder names and paths.
+scene_folder_name=$(basename "${gcs_experiment_path}")
+local_dataset_path="local_dataset"
+local_experiment_path="exp"
+DATASET_PATH="$local_experiment_path/$scene_folder_name/data"
+EXPERIMENT=$scene_folder_name
+
+# Create necessary directories.
+create_dir_if_not_exists "$local_dataset_path"
+create_dir_if_not_exists "$local_experiment_path"
+create_dir_if_not_exists "$local_experiment_path/$scene_folder_name"
+
+# Copy experiment from GCS bucket to local.
+gsutil -m cp -r "${gcs_experiment_path}/data" "$local_experiment_path/$scene_folder_name" || exit 1
+
+echo "GCS Experiment: $gcs_experiment_path"
+echo "Gin Config File: $gin_config_file"
+echo "Factor: $factor"
+echo "Scene: $scene_folder_name"
+echo "Local Dataset: $DATASET_PATH"
+echo "Local Experiment: $EXPERIMENT"
+
+accelerate launch train.py --gin_configs="$gin_config_file" \
+  --gin_bindings="Config.data_dir = '${DATASET_PATH}'" \
+  --gin_bindings="Config.exp_name = '${EXPERIMENT}'" \
+  --gin_bindings="Config.factor = ${factor}" \
+  --gin_bindings="Config.max_steps = ${max_training_steps}"
+
+gsutil -m rm -r "${gcs_experiment_path}/checkpoints/${training_job_name}"
+gsutil -m cp -r "$local_experiment_path/$scene_folder_name/config.gin" "${gcs_experiment_path}/${training_job_name}_config.gin"
+gsutil -m cp -r "$local_experiment_path/$scene_folder_name/checkpoints/*/*" "${gcs_experiment_path}/checkpoints/${training_job_name}"
@@ -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
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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
@@ -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"]
@@ -0,0 +1,90 @@
+# Dockerfile for Detectron2 serving.
+#
+# To build:
+# docker build -f model_oss/detectron2/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}
+FROM pytorch/torchserve:0.7.0-cpu
+
+USER root
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# run and update some basic packages software packages, including security libs
+RUN apt-get update &&  apt-get install -y \
+    software-properties-common && \
+    add-apt-repository -y ppa:ubuntu-toolchain-r/test && \
+    apt-get update && apt-get install -y \
+    gcc-9 g++-9 apt-transport-https ca-certificates gnupg curl
+
+# Install gcloud tools for gsutil as well as debugging
+RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | \
+    tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | \
+    apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
+    apt-get update -y && apt-get install google-cloud-sdk -y
+
+USER model-server
+
+# install detectron2 dependencies
+RUN python3 -m pip install --upgrade pip
+RUN python3 -m pip install --user numpy==1.24.2
+RUN python3 -m pip install --user opencv-python==4.7.0.72
+RUN python3 -m pip install --user 'git+https://github.com/facebookresearch/detectron2.git@v0.6'
+
+# Install GCS storage library.
+RUN pip install google-cloud-storage==2.6.0
+
+# For mask encoding.
+RUN pip install --upgrade pycocotools==2.0.6
+
+ARG MODEL_NAME=detectron2_serving
+ENV MODEL_NAME="${MODEL_NAME}"
+
+# health and prediction listener ports
+ARG AIP_HTTP_PORT=7080
+ENV AIP_HTTP_PORT="${AIP_HTTP_PORT}"
+
+ARG MODEL_MGMT_PORT=7081
+
+# expose health and prediction listener ports from the image
+EXPOSE "${AIP_HTTP_PORT}"
+EXPOSE "${MODEL_MGMT_PORT}"
+EXPOSE 8080 8081 8082 7070 7071
+
+# create torchserve configuration file
+USER root
+RUN echo "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${AIP_HTTP_PORT}\n"  \
+    "management_address=http://0.0.0.0:${MODEL_MGMT_PORT}" >> /home/model-server/config.properties
+USER model-server
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY ./model_oss/detectron2/handler.py /home/model-server/handler.py
+WORKDIR /home/model-server/
+
+# Create model archive file packaging model artifacts and dependencies.
+# Note(lavrai): The model `.pth` file and `cfg.yaml` file will be set by the
+# customer as an environment variable and will be later loaded by the
+# `handler.py` file.
+RUN torch-model-archiver \
+  --model-name="${MODEL_NAME}" \
+  --version=1.0 \
+  --handler=/home/model-server/handler.py \
+  --export-path=/home/model-server/model-store \
+  -f
+
+# run Torchserve HTTP serve to respond to prediction requests
+CMD ["ls", "-ltr", "/home/model-server/model-store/", ";", \
+    "torchserve", "--start", "--ts-config=/home/model-server/config.properties", \
+    "--models", "${MODEL_NAME}=${MODEL_NAME}.mar", \
+    "--model-store", "/home/model-server/model-store"]
@@ -0,0 +1,100 @@
+# Dockerfile for Detectron2 training.
+#
+# To build:
+# docker build -f model_oss/detectron2/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 nvidia/cuda:11.1.1-cudnn8-devel-ubuntu18.04
+# Using an older system (18.04) to avoid opencv incompatibility (issue#3524).
+
+ENV DEBIAN_FRONTEND noninteractive
+RUN apt-get update && apt-get install -y \
+    python3.7 python3.7-dev python3.7-distutils \
+    python3-opencv ca-certificates git wget sudo ninja-build \
+    curl wget vim
+
+# Make python3 available for python3.7.
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.6 1
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.7 2
+RUN update-alternatives --config python3
+# Make python available for python3.7.
+RUN ln -sv /usr/bin/python3.7 /usr/bin/python
+
+# Create a non-root user.
+ARG USER_ID=1000
+RUN useradd -m --no-log-init --system  --uid ${USER_ID} appuser -g sudo
+RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
+USER appuser
+WORKDIR /home/appuser
+
+ENV PATH="/home/appuser/.local/bin:${PATH}"
+RUN wget https://bootstrap.pypa.io/pip/get-pip.py && \
+    python3.7 get-pip.py --user && \
+    rm get-pip.py
+
+# Important! Otherwise, it uses existing numpy from host-modules
+# which throws error.
+RUN pip install --user numpy==1.20.3
+
+# Install dependencies:
+# See https://pytorch.org/ for other options if you use
+# a different version of CUDA.
+RUN pip install --user tensorboard==2.11.0
+# cmake from apt-get is too old.
+RUN pip install --user cmake==3.25.2
+RUN pip install --user torch==1.10.0+cu111 torchvision==0.11.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
+RUN pip install --user setuptools==59.5.0
+RUN pip install --user opencv-python==4.7.0.72
+RUN pip install --user cloudml-hypertune==0.1.0.dev6
+RUN pip install --user fvcore==0.1.5.post20221221
+# Install detectron2.
+RUN git clone -b v0.6 https://github.com/facebookresearch/detectron2 detectron2_repo
+# Set FORCE_CUDA because during `docker build` cuda is not accessible.
+ENV FORCE_CUDA="1"
+# This will by default build detectron2 for all common cuda
+# architectures and take a lot more time,
+# because inside `docker build`, there is no way to tell
+# which architecture will be used.
+ARG TORCH_CUDA_ARCH_LIST="Kepler;Kepler+Tesla;Maxwell;Maxwell+Tegra;Pascal;Volta;Turing"
+ENV TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}"
+RUN pip install --user -e detectron2_repo
+
+# Set a fixed model cache directory.
+ENV FVCORE_CACHE="/tmp"
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model-garden detectron2 files to '/home/appuser/trainer' folder.
+ADD ./model_oss/detectron2 /home/appuser/trainer
+
+################ Copy plain_train_net.py to task.py and
+# then modify it using sed commands. ###################
+# Src: https://github.com/facebookresearch/detectron2/blob/v0.6/tools/plain_train_net.py
+RUN sudo cp /home/appuser/detectron2_repo/tools/plain_train_net.py /home/appuser/trainer/task.py
+# Make additional changes to task.py.
+# Note(lavrai): Start adding SED commands from end of file towards the top
+# so that the line numbers do not keep changing for the source file.
+# For entry-point:
+RUN sudo sed -i "214 d" /home/appuser/trainer/task.py
+RUN sudo sed -i "213 a\    default_arg_parser = default_argument_parser()" /home/appuser/trainer/task.py
+RUN sudo sed -i "214 a\    extended_parser = trainer_utils.extend_parser_arguments(default_arg_parser)" /home/appuser/trainer/task.py
+RUN sudo sed -i "215 a\    args = extended_parser.parse_args()" /home/appuser/trainer/task.py
+# For main() function:
+RUN sudo sed -i "192 a\    trainer_utils.register_dataset(args)" /home/appuser/trainer/task.py
+# For setup() function:
+RUN sudo sed -i "184 a\    cfg.SOLVER.BASE_LR = args.lr" /home/appuser/trainer/task.py
+RUN sudo sed -i "185 a\    cfg.OUTPUT_DIR = args.output_dir" /home/appuser/trainer/task.py
+RUN sudo sed -i "186 a\    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(config_file_copy)" /home/appuser/trainer/task.py
+RUN sudo sed -i "182 a\    config_file_copy = args.config_file" /home/appuser/trainer/task.py
+RUN sudo sed -i "183 a\    args.config_file = model_zoo.get_config_file(args.config_file)" /home/appuser/trainer/task.py
+# For new import:
+RUN sudo sed -i "27 a\from detectron2 import model_zoo" /home/appuser/trainer/task.py
+RUN sudo sed -i "21 a\import trainer_utils" /home/appuser/trainer/task.py
+
+ENV PYTHONPATH /home/appuser/trainer
+
+ENTRYPOINT ["python", "-m", "trainer.task"]
@@ -0,0 +1,154 @@
+"""Custom handler for Detectron2 serving."""
+
+import io
+import json
+import os
+from typing import Any, List, Tuple
+
+import cv2
+from detectron2.config import get_cfg
+from detectron2.engine import DefaultPredictor
+from google.cloud import storage
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+
+
+def get_bucket_and_blob_name(gcs_filepath: str) -> Tuple[str, str]:
+  """Gets bucket and blob name from gcs path."""
+  # The gcs path is of the form gs://<bucket-name>/<blob-name>
+  gs_suffix = gcs_filepath.split("gs://", 1)[1]
+  return tuple(gs_suffix.split("/", 1))
+
+
+def download_gcs_file(src_file_path: str, dst_file_path: str):
+  """Downloads gcs-file to local folder."""
+  src_bucket_name, src_blob_name = get_bucket_and_blob_name(src_file_path)
+  client = storage.Client()
+  src_bucket = client.get_bucket(src_bucket_name)
+  src_blob = src_bucket.blob(src_blob_name)
+  src_blob.download_to_filename(dst_file_path)
+
+
+class ModelHandler:
+  """Custom model handler for Detectron2."""
+
+  def __init__(self):
+    self.error = None
+    self._batch_size = 0
+    self.initialized = False
+    self.predictor = None
+    self.test_threshold = 0.5
+
+  def initialize(self, context: Any):
+    """Initialize."""
+    print("context.system_properties: ", context.system_properties)
+    print("context.manifest: ", context.manifest)
+    self.manifest = context.manifest
+    properties = context.system_properties
+    # Get threshold from environment variable.
+    # This will be set by customer.
+    self.test_threshold = float(os.environ.get("TEST_THRESHOLD"))
+    print("test_threshold: ", self.test_threshold)
+    # Get model and config file location from environment variables.
+    # These will be set by customer when doing model upload.
+    gcs_model_file = os.environ["MODEL_PTH_FILE"]
+    gcs_config_file = os.environ["CONFIG_YAML_FILE"]
+    print("Copying gcs_model_file: ", gcs_model_file)
+    print("Copying gcs_config_file: ", gcs_config_file)
+    # Copy these files from GCS location to local file.
+    # Note(lavrai): GCSFuse path does not seem to work here for now.
+    model_file = "./model.pth"
+    config_file = "./cfg.yaml"
+    download_gcs_file(src_file_path=gcs_model_file, dst_file_path=model_file)
+    if not os.path.exists(model_file):
+      raise RuntimeError("Missing model_file: %s" % model_file)
+    download_gcs_file(src_file_path=gcs_config_file, dst_file_path=config_file)
+    if not os.path.exists(config_file):
+      raise RuntimeError("Missing config_file: %s" % config_file)
+
+    # Set up config file.
+    cfg = get_cfg()
+    cfg.merge_from_file(config_file)
+    cfg.MODEL.WEIGHTS = model_file
+    cfg.MODEL.DEVICE = (
+        cfg.MODEL.DEVICE + str(properties.get("gpu_id"))
+        if torch.cuda.is_available()
+        else "cpu"
+    )
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = self.test_threshold
+
+    # Build predictor from config.
+    self.predictor = DefaultPredictor(cfg)
+    self._batch_size = context.system_properties["batch_size"]
+    self.initialized = True
+
+  def preprocess(self, batch: List[Any]) -> List[Any]:
+    """Preprocess raw input and return as list of images."""
+    print("Running pre-processing.")
+    images = []
+    for request in batch:
+      request_data = request.get("data")
+      input_bytes = io.BytesIO(request_data)
+      img = cv2.imdecode(np.fromstring(input_bytes.read(), np.uint8), 1)
+      images.append(img)
+    return images
+
+  def inference(self, model_input: List[Any]) -> List[Any]:
+    """Runs inference."""
+    print("Running model-inference.")
+    return [self.predictor(image) for image in model_input]
+
+  def postprocess(self, inference_result: List[Any]) -> List[Any]:
+    """Post process inference result."""
+    response_list = []
+    print("Num inference_items are:", len(inference_result))
+    for inference_item in inference_result:
+      predictions = inference_item["instances"].to("cpu")
+      print("Predictions are:", predictions)
+      boxes = None
+      if predictions.has("pred_boxes"):
+        boxes = predictions.pred_boxes.tensor.numpy().tolist()
+      scores = None
+      if predictions.has("scores"):
+        scores = predictions.scores.numpy().tolist()
+      classes = None
+      if predictions.has("pred_classes"):
+        classes = predictions.pred_classes.numpy().tolist()
+      masks_rle = None
+      if predictions.has("pred_masks"):
+        # Do run length encoding, else the mask output becomes huge.
+        masks_rle = [
+            mask_util.encode(np.asfortranarray(mask))
+            for mask in predictions.pred_masks
+        ]
+        for rle in masks_rle:
+          rle["counts"] = rle["counts"].decode("utf-8")
+      response = {
+          "classes": classes,
+          "scores": scores,
+          "boxes": boxes,
+          "masks_rle": masks_rle,
+      }
+      response_list.append(json.dumps(response))
+    print("response_list: ", response_list)
+    return response_list
+
+  def handle(self, data: Any, context: Any) -> List[Any]:  # pylint: disable=unused-argument
+    """Runs preprocess, inference, and post-processing."""
+    model_input = self.preprocess(data)
+    model_out = self.inference(model_input)
+    output = self.postprocess(model_out)
+    print("Done handling input.")
+    return output
+
+
+_service = ModelHandler()
+
+
+def handle(data: Any, context: Any) -> List[Any]:
+  if not _service.initialized:
+    _service.initialize(context)
+  if data is None:
+    return None
+  return _service.handle(data, context)
@@ -0,0 +1,97 @@
+"""Detectron2 trainer helper functions."""
+
+import argparse
+from detectron2.data.datasets import register_coco_instances
+
+
+def extend_parser_arguments(
+    parser: argparse.ArgumentParser,
+) -> argparse.ArgumentParser:
+  """Adds additional model-garden related arguments."""
+  parser.add_argument(
+      "--train_dataset_name",
+      required=False,
+      default="",
+      type=str,
+      help=(
+          "The training dataset name for registration. "
+          "For example: 'balloon_train'."
+      ),
+  )
+  parser.add_argument(
+      "--train_coco_json_file",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the training coco-json format file.",
+  )
+  parser.add_argument(
+      "--train_image_root",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the root folder containing the training images.",
+  )
+  parser.add_argument(
+      "--val_dataset_name",
+      required=False,
+      default="",
+      type=str,
+      help=(
+          "The validation dataset name for registration. "
+          "For example: 'balloon_val'."
+      ),
+  )
+  parser.add_argument(
+      "--val_coco_json_file",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the validation coco-json format file.",
+  )
+  parser.add_argument(
+      "--val_image_root",
+      required=False,
+      default="",
+      type=str,
+      help="The path to the root folder containing the validation images.",
+  )
+  parser.add_argument(
+      "--output_dir",
+      required=True,
+      type=str,
+      help="The path to the output directory.",
+  )
+  # Add hyper-parameter tuning related variables.
+  parser.add_argument(
+      "--lr",
+      type=float,
+      default=0.00025,
+      help="The learning rate to be tuned.",
+  )
+  parser.add_argument(
+      "--hp_eval_task",
+      type=str,
+      choices=["bbox", "segm"],
+      default="bbox",
+      help="The task choice for HP tuning.",
+  )
+  return parser
+
+
+def register_dataset(args: argparse.Namespace):
+  """Register the input dataset in Detectron2 Coco format."""
+  if args.train_dataset_name:
+    register_coco_instances(
+        name=args.train_dataset_name,
+        metadata={},
+        json_file=args.train_coco_json_file,
+        image_root=args.train_image_root,
+    )
+  if args.val_dataset_name:
+    register_coco_instances(
+        name=args.val_dataset_name,
+        metadata={},
+        json_file=args.val_coco_json_file,
+        image_root=args.val_image_root,
+    )
@@ -0,0 +1,74 @@
+# Dockerfile for Diffuser Serving.
+#
+# To build:
+# docker build -f model_oss/diffusers/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="diffusers_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install torch==1.13.1
+RUN pip install torchvision==0.14.1
+RUN pip install transformers==4.27.4
+RUN pip install datasets==2.9.0
+RUN pip install accelerate==0.17.0
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.16
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install imageio[ffmpeg]==2.31.0
+RUN pip install absl-py==1.4.0
+
+# Copy LICENSE file
+RUN apt-get update && apt-get install wget
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install diffusers from main branch source code with a pinned commit.
+RUN git clone --depth 1 --branch v0.18.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+
+# Copy model artifacts.
+COPY model_oss/diffusers/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"]
@@ -0,0 +1,47 @@
+# Dockerfile for Diffuser Training.
+#
+# To build:
+# docker build -f model_oss/diffusers/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 \
+        git \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install libraries.
+RUN pip install torchvision==0.14.1
+RUN pip install transformers==4.26.1
+RUN pip install datasets==2.9.0
+RUN pip install accelerate==0.17.0
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.16
+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
+
+# Install diffusers from main branch source code with a pinned commit.
+RUN git clone --depth 1 --branch v0.18.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+
+# Switch to diffusers examples folder.
+WORKDIR examples
+
+# Config accelerate.
+COPY model_oss/diffusers/train.sh train.sh
+
+# Generate accelerate config at the beginning of docker run.
+ENTRYPOINT ["/bin/bash", "train.sh"]
@@ -0,0 +1,256 @@
+"""Custom handler for huggingface/diffusers models."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=logging-fstring-interpolation
+
+import base64
+import io
+import logging
+import os
+from typing import Any, List, Sequence, Tuple
+
+from diffusers import ControlNetModel
+from diffusers import DiffusionPipeline
+from diffusers import DPMSolverMultistepScheduler
+from diffusers import EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionControlNetPipeline
+from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionInstructPix2PixPipeline
+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
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import constants
+from util import fileutils
+from util import image_format_converter
+from video_util import video_format_converter
+
+STABLE_DIFFUSION_MODEL = "runwayml/stable-diffusion-v1-5"
+
+# Tasks
+TEXT_TO_IMAGE = "text-to-image"
+IMAGE_TO_IMAGE = "image-to-image"
+IMAGE_INPAINTING = "image-inpainting"
+INSTRUCT_PIX2PIX = "instruct-pix2pix"
+CONTROLNET = "controlnet"
+CONDITIONED_SUPER_RES = "conditioned-super-res"
+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."""
+
+  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
+
+    self.model_id = os.environ["MODEL_ID"]
+    if self.model_id.startswith(constants.GCS_URI_PREFIX):
+      gcs_path = self.model_id[len(constants.GCS_URI_PREFIX) :]
+      local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
+      logging.info(f"Download {self.model_id} to {local_model_dir}")
+      fileutils.download_gcs_dir_to_local(self.model_id, local_model_dir)
+      self.model_id = local_model_dir
+
+    self.task = os.environ.get("TASK", TEXT_TO_IMAGE)
+    logging.info(f"Using task:{self.task}, model:{self.model_id}")
+
+    if self.task == TEXT_TO_IMAGE:
+      pipeline = StableDiffusionPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == IMAGE_TO_IMAGE:
+      pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == IMAGE_INPAINTING:
+      pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == INSTRUCT_PIX2PIX:
+      pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == CONTROLNET:
+      controlnet = ControlNetModel.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+          STABLE_DIFFUSION_MODEL,
+          controlnet=controlnet,
+          torch_dtype=torch.float16,
+      )
+      pipeline.scheduler = UniPCMultistepScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline.enable_xformers_memory_efficient_attention()
+      pipeline.enable_model_cpu_offload()
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == CONDITIONED_SUPER_RES:
+      pipeline = StableDiffusionUpscalePipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      # This is necessary to 4x upscale >=256x256 input images with V100.
+      logging.info("Enable xformers memory efficient attention for inference.")
+      pipeline.enable_xformers_memory_efficient_attention()
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == TEXT_TO_VIDEO_ZERO_SHOT:
+      pipeline = TextToVideoZeroPipeline.from_pretrained(
+          STABLE_DIFFUSION_MODEL, torch_dtype=torch.float16
+      )
+      # Memory optimization.
+      pipeline.enable_xformers_memory_efficient_attention()
+      pipeline.enable_model_cpu_offload()
+      pipeline = pipeline.to(self.map_location)
+    elif self.task == TEXT_TO_VIDEO:
+      pipeline = DiffusionPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16, variant="fp16"
+      )
+      pipeline.enable_model_cpu_offload()
+      # Memory optimization.
+      pipeline.enable_vae_slicing()
+      pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+          pipeline.scheduler.config
+      )
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+
+    self.pipeline = pipeline
+    self.initialized = True
+    logging.info("Handler initialization done.")
+
+  def preprocess(self, data: Any) -> Tuple[Any, Any, Any]:
+    """Preprocess input data."""
+    prompts = [item["prompt"] for item in data]
+    images = None
+    mask_images = None
+
+    if "image" in data[0]:
+      images = [
+          image_format_converter.base64_to_image(item["image"]) for item in data
+      ]
+    if "mask_image" in data[0]:
+      mask_images = [
+          image_format_converter.base64_to_image(item["mask_image"])
+          for item in data
+      ]
+    return prompts, images, mask_images
+
+  def inference(self, data: Any, *args, **kwargs) -> List[Image.Image]:
+    """Run the inference."""
+    prompts, images, mask_images = data
+    if self.task == TEXT_TO_IMAGE:
+      predicted_images = self.pipeline(prompt=prompts).images
+    elif self.task == IMAGE_TO_IMAGE:
+      predicted_images = self.pipeline(prompt=prompts, image=images).images
+    elif self.task == IMAGE_INPAINTING:
+      predicted_images = self.pipeline(
+          prompt=prompts, image=images, mask_image=mask_images
+      ).images
+    elif self.task == INSTRUCT_PIX2PIX:
+      predicted_images = self.pipeline(prompt=prompts, image=images).images
+    elif self.task == CONTROLNET:
+      predicted_images = self.pipeline(
+          prompt=prompts, image=images, num_inference_steps=20
+      ).images
+    elif self.task == CONDITIONED_SUPER_RES:
+      predicted_images = self.pipeline(
+          prompt=prompts, image=images, num_inference_steps=20
+      ).images
+    elif self.task == TEXT_TO_VIDEO_ZERO_SHOT:
+      # For each given prompt, generate a short video.
+      # The pipeline doesn't support multiple prompts in one run yet.
+      videos = []
+      for prompt in prompts:
+        numpy_arrays = self.pipeline(prompt=prompt).images
+        numpy_arrays = [(i * 255).astype("uint8") for i in numpy_arrays]
+        videos.append(
+            frames_to_video_bytes(numpy_arrays, fps=4)
+        )
+      return videos
+    elif self.task == TEXT_TO_VIDEO:
+      predicted_images = np.asarray(self.pipeline(prompt=prompts).frames)
+      # For multiple prompts, the model concatenates video frames, i.e. the
+      # output shape is (num_frames, height, width * len(prompts), channels).
+      # Therefore we need to split the output into different videos.
+      predicted_images = np.array_split(predicted_images, len(prompts), axis=2)
+      videos = [
+          frames_to_video_bytes(images, fps=8)
+          for images in predicted_images
+      ]
+      return videos
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+    return predicted_images
+
+  def postprocess(self, data: Any) -> List[str]:
+    """Convert the images to base64 string."""
+    outputs = []
+    for prediction in data:
+      if isinstance(prediction, bytes):
+        # This is the video bytes.
+        outputs.append(base64.b64encode(prediction).decode("utf-8"))
+      else:
+        outputs.append(image_format_converter.image_to_base64(prediction))
+    return outputs
+
+
+# pylint: enable=logging-fstring-interpolation
@@ -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 "$@"
@@ -0,0 +1,83 @@
+# This Dockerfile converts JAX vision transformer model to
+# tensorflow saved model format.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="jax-f-vlm-model-conversion:${USER}-test"
+# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_model_conversion.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+
+# See https://cloud.google.com/tensorflow-enterprise/docs/overview for details.
+FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-12.py310:m110
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# 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 \
+        libgl1
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install required libs
+RUN pip install --upgrade pip
+
+# Using the commit 6712c224985c694001ba8ee68697bbf4dcb32edb on Jan 4th, 2024.
+ARG COMMIT_ID=6712c224985c694001ba8ee68697bbf4dcb32edb
+RUN git clone -c \
+remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
+https://github.com/google-research/google-research --no-checkout --progress \
+--depth 1
+WORKDIR ./google-research
+RUN git sparse-checkout init --cone
+RUN git sparse-checkout set fvlm
+RUN git checkout ${COMMIT_ID}
+
+# The following pip installs are pinned down versions satisfying
+# fvlm/requirements.txt file.
+# NOTE: Using `no-deps` flag to avoid overwriting of dependent library
+# versions. For example, both `chex` and `jax` can overwrite each other's
+# `jax-lib` version.
+# Note: The following libraries are pinned down versions of:
+# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
+RUN pip install --no-cache-dir tensorflow==2.12.0
+RUN pip install --no-cache-dir tensorflow-datasets==4.9.2
+RUN pip install --no-cache-dir numpy==1.23.5
+RUN pip install --no-cache-dir torch==2.0.1
+RUN pip install --no-cache-dir torchvision==0.15.2
+RUN pip install --no-cache-dir opencv-python==4.7.0.72
+RUN pip install --no-cache-dir tqdm==4.65.0
+RUN pip install --no-cache-dir git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33
+RUN pip install --no-cache-dir Pillow==9.5.0
+RUN pip install --no-cache-dir orbax-checkpoint==0.3.3
+RUN pip install --no-cache-dir gin-config==0.5.0
+RUN pip install --no-cache-dir pycocotools==2.0.6
+RUN pip install --no-cache-dir contextlib2==21.6.0
+RUN pip install --no-cache-dir ml-collections==0.1.1
+RUN pip install --no-cache-dir chex==0.1.7
+RUN pip install --no-cache-dir optax==0.1.5
+# Dependencies already included. Use no-deps to not update numpy.
+RUN pip install --no-cache-dir --no-deps flax==0.7.2
+RUN pip install --no-cache-dir --no-deps clu==0.0.9
+RUN pip install --no-cache-dir jax[cuda11_cudnn86]==0.4.9 \
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+RUN pip install --no-cache-dir ml-dtypes==0.2.0
+RUN pip install --no-cache-dir tensorflow_text==2.12.0
+
+WORKDIR ./fvlm
+ENV PYTHONPATH ./
+
+ENTRYPOINT ["python", "export_saved_model.py"]
@@ -0,0 +1,78 @@
+# This Dockerfile trains the F-VLM model on GPU.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="jax-f-vlm-train:${USER}-test"
+# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_train_gpu.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+# See https://cloud.google.com/tensorflow-enterprise/docs/overview for details.
+FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-12.py310:m110
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# 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
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install required libs
+RUN pip install --upgrade pip
+# The following pip installs are pinned down versions satisfying
+# fvlm/requirements.txt file.
+# Get F-VLM repository by using git sparse-checkout to avoid downloading entire
+# google-research repository.
+# Using the commit 6712c224985c694001ba8ee68697bbf4dcb32edb on Jan 4th, 2024.
+ARG COMMIT_ID=6712c224985c694001ba8ee68697bbf4dcb32edb
+RUN git clone -c \
+remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
+https://github.com/google-research/google-research --no-checkout --progress \
+--depth 1
+WORKDIR ./google-research
+RUN git sparse-checkout init --cone
+RUN git sparse-checkout set fvlm
+RUN git checkout ${COMMIT_ID}
+
+# Note: The following libraries are pinned down versions of:
+# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
+RUN pip install --no-cache-dir tensorflow==2.12.0
+RUN pip install --no-cache-dir tensorflow-datasets==4.9.2
+RUN pip install --no-cache-dir numpy==1.23.5
+RUN pip install --no-cache-dir torch==2.0.1
+RUN pip install --no-cache-dir torchvision==0.15.2
+RUN pip install --no-cache-dir opencv-python==4.7.0.72
+RUN pip install --no-cache-dir tqdm==4.65.0
+RUN pip install --no-cache-dir git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33
+RUN pip install --no-cache-dir Pillow==9.5.0
+RUN pip install --no-cache-dir orbax-checkpoint==0.3.3
+RUN pip install --no-cache-dir gin-config==0.5.0
+RUN pip install --no-cache-dir pycocotools==2.0.6
+RUN pip install --no-cache-dir contextlib2==21.6.0
+RUN pip install --no-cache-dir ml-collections==0.1.1
+RUN pip install --no-cache-dir chex==0.1.7
+RUN pip install --no-cache-dir optax==0.1.5
+# Dependencies already included. Use no-deps to not update numpy.
+RUN pip install --no-cache-dir --no-deps flax==0.7.2
+RUN pip install --no-cache-dir --no-deps clu==0.0.9
+# Installing jax at the very end with GPU support.
+# NOTE: Not using `no-deps` flag here because we need CUDA support.
+RUN pip install --no-cache-dir jax[cuda11_cudnn86]==0.4.9 \
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+WORKDIR ./fvlm
+ENV PYTHONPATH ./
+
+ENTRYPOINT ["python", "train_and_eval.py"]
@@ -0,0 +1,138 @@
+# This Dockerfile trains the F-VLM model on TPU.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="jax-f-vlm-train-tpu:${USER}-test"
+# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_train_tpu.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+FROM python:3.11
+
+# Get libtpu shared library. See go/what-is-libtpu.
+RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# 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 \
+        libgl1
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install required libs
+RUN pip install --upgrade pip
+
+# Get F-VLM repository by using git sparse-checkout to avoid downloading entire
+# google-research repository.
+# Using the commit 05ece4b1c97285b48b51fa44321ccb2cb347406a on Dec 11th, 2023.
+ARG COMMIT_ID=05ece4b1c97285b48b51fa44321ccb2cb347406a
+RUN git clone -c \
+remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
+https://github.com/google-research/google-research --no-checkout --progress \
+--depth 1
+WORKDIR ./google-research
+RUN git sparse-checkout init --cone
+RUN git sparse-checkout set fvlm
+RUN git checkout ${COMMIT_ID}
+
+# Note: The following libraries are pinned down versions of:
+# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
+RUN pip install --no-cache-dir ml_dtypes==0.3.1
+RUN pip install --no-cache-dir tensorstore==0.1.51
+RUN pip install --no-cache-dir MarkupSafe==2.1.3
+RUN pip install --no-cache-dir Pillow==9.5.0
+RUN pip install --no-cache-dir PyYAML==6.0.1
+RUN pip install --no-cache-dir absl_py==1.4.0
+RUN pip install --no-cache-dir array_record==0.4.1
+RUN pip install --no-cache-dir astunparse==1.6.3
+RUN pip install --no-cache-dir cachetools==5.3.1
+RUN pip install --no-cache-dir certifi==2023.7.22
+RUN pip install --no-cache-dir charset_normalizer==3.3.0
+RUN pip install --no-cache-dir chex==0.1.83
+RUN pip install --no-cache-dir click==8.1.7
+RUN pip install --no-cache-dir clip==0.2.0
+RUN pip install --no-cache-dir clu==0.0.9
+RUN pip install --no-cache-dir contourpy==1.1.1
+RUN pip install --no-cache-dir cycler==0.12.1
+RUN pip install --no-cache-dir dm_tree==0.1.8
+RUN pip install --no-cache-dir etils==1.5.1
+RUN pip install --no-cache-dir filelock==3.12.4
+RUN pip install --no-cache-dir flatbuffers==23.5.26
+RUN pip install --no-cache-dir flax==0.7.4
+RUN pip install --no-cache-dir fonttools==4.43.1
+RUN pip install --no-cache-dir fsspec==2023.9.2
+RUN pip install --no-cache-dir ftfy==6.1.1
+RUN pip install --no-cache-dir gast==0.5.4
+RUN pip install --no-cache-dir gin_config==0.5.0
+RUN pip install --no-cache-dir google_auth==2.23.3
+RUN pip install --no-cache-dir google_auth_oauthlib==1.0.0
+RUN pip install --no-cache-dir google_pasta==0.2.0
+RUN pip install --no-cache-dir googleapis_common_protos==1.61.0
+RUN pip install --no-cache-dir grpcio==1.59.0
+RUN pip install --no-cache-dir h5py==3.10.0
+RUN pip install --no-cache-dir importlib_resources==6.1.0
+RUN pip install --no-cache-dir 'jax[tpu]==0.4.18' \
+-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+RUN pip install --no-cache-dir jaxlib==0.4.18
+RUN pip install --no-cache-dir jinja2==3.1.2
+RUN pip install --no-cache-dir keras==2.14.0
+RUN pip install --no-cache-dir kiwisolver==1.4.5
+RUN pip install --no-cache-dir libclang==16.0.6
+RUN pip install --no-cache-dir markdown==3.5
+RUN pip install --no-cache-dir matplotlib==3.8.0
+RUN pip install --no-cache-dir mpmath==1.3.0
+RUN pip install --no-cache-dir networkx==3.1
+RUN pip install --no-cache-dir numpy==1.26.0
+RUN pip install --no-cache-dir nvidia_cublas_cu12==12.1.3.1
+RUN pip install --no-cache-dir nvidia_cuda_cupti_cu12==12.1.105
+RUN pip install --no-cache-dir nvidia_cuda_nvrtc_cu12==12.1.105
+RUN pip install --no-cache-dir nvidia_cuda_runtime_cu12==12.1.105
+RUN pip install --no-cache-dir nvidia_cudnn_cu12==8.9.2.26
+RUN pip install --no-cache-dir nvidia_cufft_cu12==11.0.2.54
+RUN pip install --no-cache-dir nvidia_curand_cu12==10.3.2.106
+RUN pip install --no-cache-dir nvidia_cusolver_cu12==11.4.5.107
+RUN pip install --no-cache-dir nvidia_cusparse_cu12==12.1.0.106
+RUN pip install --no-cache-dir nvidia_nccl_cu12==2.18.1
+RUN pip install --no-cache-dir nvidia_nvjitlink_cu12==12.2.140
+RUN pip install --no-cache-dir nvidia_nvtx_cu12==12.1.105
+RUN pip install --no-cache-dir opencv_python==4.8.1.78
+RUN pip install --no-cache-dir orbax_checkpoint==0.4.1
+RUN pip install --no-cache-dir promise==2.3
+RUN pip install --no-cache-dir protobuf==3.20.3
+RUN pip install --no-cache-dir psutil==5.9.5
+RUN pip install --no-cache-dir pyasn1==0.5.0
+RUN pip install --no-cache-dir pycocotools==2.0.7
+RUN pip install --no-cache-dir pygments==2.16.1
+RUN pip install --no-cache-dir regex==2023.10.3
+RUN pip install --no-cache-dir rich==13.6.0
+RUN pip install --no-cache-dir scipy==1.11.3
+RUN pip install --no-cache-dir sympy==1.12
+RUN pip install --no-cache-dir tensorboard==2.14.1
+RUN pip install --no-cache-dir tensorboard_data_server==0.7.1
+RUN pip install --no-cache-dir tensorflow==2.14.0
+RUN pip install --no-cache-dir tensorflow_datasets==4.9.3
+RUN pip install --no-cache-dir torch==2.1.0
+RUN pip install --no-cache-dir torchvision==0.16.0
+RUN pip install --no-cache-dir urllib3==2.0.6
+RUN pip install --no-cache-dir wcwidth==0.2.8
+RUN pip install --no-cache-dir werkzeug==3.0.0
+RUN pip install --no-cache-dir wheel==0.41.2
+RUN pip install --no-cache-dir tensorflow_text==2.14.0
+
+WORKDIR ./fvlm
+ENV PYTHONPATH ./
+
+ENTRYPOINT ["python", "train_and_eval.py"]
@@ -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={\
+  "imagebind_serving": {\
+    "1.0": {\
+        "defaultVersion": true,\
+        "marName": "imagebind_serving.mar",\
+        "minWorkers": 1,\
+        "maxWorkers": 1,\
+        "batchSize": 1\
+    }\
+  }\
+}
@@ -0,0 +1,74 @@
+# Dockerfile for the serving docker for ImageBind.
+#
+# To build:
+# docker build -f model_oss/imagebind/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="imagebind_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Install tools.
+RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        git \
+        libgeos-dev
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install absl-py==1.4.0
+RUN pip install google-cloud-storage==2.7.0
+
+# Install ImageBind and dependencies.
+RUN git clone https://github.com/facebookresearch/ImageBind.git
+WORKDIR ImageBind
+# Pin the commit at 07/14/2023.
+RUN git reset --hard 95d27c7fd5a8362f3527e176c3a80ae5a4d880c0
+# Modify tokenizer file path from ImageBind repo to work with the server.
+RUN sed -i '25d' imagebind/data.py
+RUN sed -i '25 i\BPE_PATH = "/home/model-server/ImageBind/bpe/bpe_simple_vocab_16e6.txt.gz"' imagebind/data.py
+RUN pip install .
+
+WORKDIR /home/model-server
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY model_oss/imagebind/handler.py /home/model-server/handler.py
+COPY model_oss/imagebind/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}
+
+# 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"]
@@ -0,0 +1,277 @@
+"""Custom handler for the ImageBind model."""
+
+import logging
+import os
+from typing import Any, Dict, List
+
+from imagebind import data as data_util
+from imagebind.models import imagebind_model
+from imagebind.models.imagebind_model import ModalityType
+from PIL import Image
+import torch
+from torchvision import transforms
+from ts.torch_handler import base_handler
+
+from util import constants
+from util import fileutils
+
+
+_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE = "video"
+
+
+class ImageBindHandler(base_handler.BaseHandler):
+  """Custom handler for the ImageBind model.
+
+  Attributes:
+    map_location: Mapping storage location.
+    device: Device on which to run inference.
+    manifest: TorchServe manifest.
+    task: Task for which to run the ImageBind model.
+    model: ImageBind model instance.
+  """
+
+  def initialize(self, context: Any) -> None:
+    """Initializes the ImageBind model handler.
+
+    Args:
+      context: TorchServe context, which contains system information and the
+        manifest.
+
+    Raises:
+      ValueError: A task that is unsupported by the 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.task = os.environ.get("TASK", constants.FEATURE_EMBEDDING_GENERATION)
+    if self.task not in [
+        constants.FEATURE_EMBEDDING_GENERATION,
+        constants.ZERO_SHOT_CLASSIFICATION,
+    ]:
+      raise ValueError(f"Invalid task: {self.task}.")
+    logging.info(
+        "Handler initializing ImageBind pretrained model for task %s.",
+        self.task,
+    )
+
+    self.model = imagebind_model.imagebind_huge(pretrained=True)
+    self.model.eval()
+    self.model.to(self.device)
+
+    logging.info("Initialized ImageBind pretrained model.")
+    self.initialized = True
+
+  def preprocess(self, data: Any) -> List[Dict[str, Any]]:
+    """Preprocesses input data, including text, image, audio and video data.
+
+    Args:
+      data: Input data.
+
+    Returns:
+      A list of processed data samples, with each sample being a dictionary of
+      modality (key): input (value) pairs.
+    """
+    logging.info("Preprocessing: %d instances received.", len(data))
+    preprocessed_sample_list = []
+    for item in data:
+      preprocessed_sample = {}
+      if ModalityType.TEXT in item:
+        preprocessed_sample[ModalityType.TEXT] = (
+            data_util.load_and_transform_text(
+                item[ModalityType.TEXT], self.device
+            )
+        )
+      for image_modality in [
+          ModalityType.VISION,
+          ModalityType.DEPTH,
+          ModalityType.THERMAL,
+      ]:
+        if image_modality in item:
+          image_paths = item[image_modality]
+          local_image_paths = fileutils.download_gcs_file_list_to_local(
+              image_paths, constants.LOCAL_DATA_DIR
+          )
+          is_depth_or_thermal = image_modality in [
+              ModalityType.DEPTH,
+              ModalityType.THERMAL,
+          ]
+          preprocessed_sample[image_modality] = (
+              self._load_and_transform_image_data(
+                  local_image_paths,
+                  self.device,
+                  is_depth_or_thermal=is_depth_or_thermal,
+              )
+          )
+      if ModalityType.AUDIO in item:
+        audio_paths = item[ModalityType.AUDIO]
+        local_audio_paths = fileutils.download_gcs_file_list_to_local(
+            audio_paths, constants.LOCAL_DATA_DIR
+        )
+        preprocessed_sample[ModalityType.AUDIO] = (
+            data_util.load_and_transform_audio_data(
+                local_audio_paths, self.device
+            )
+        )
+      if _VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE in item:
+        video_paths = item[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE]
+        local_video_paths = fileutils.download_gcs_file_list_to_local(
+            video_paths, constants.LOCAL_DATA_DIR
+        )
+        preprocessed_sample[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE] = (
+            data_util.load_and_transform_video_data(
+                local_video_paths, self.device
+            )
+        )
+      if ModalityType.IMU in item:
+        # Input data in the IMU modality are expected in shape [B, 6, 2000].
+        preprocessed_sample[ModalityType.IMU] = torch.tensor(
+            item[ModalityType.IMU], dtype=torch.float32, device=self.device
+        )
+      if preprocessed_sample:
+        preprocessed_sample_list.append(preprocessed_sample)
+    return preprocessed_sample_list
+
+  def _load_and_transform_image_data(
+      self,
+      image_paths: List[str],
+      device: torch.device,
+      is_depth_or_thermal: bool = False,
+  ) -> torch.Tensor:
+    """Loads and transforms 3-channel images, depth images and thermal images.
+
+    Args:
+      image_paths: A list of image paths.
+      device: Device onto which to load images.
+      is_depth_or_thermal: Whether the images are depth or thermal images.
+
+    Returns:
+      A list of processed tensors corresponding to the input images.
+
+    Raises:
+      ValueError: The input image_paths is None.
+    """
+    if image_paths is None:
+      raise ValueError("image_paths must not be None.")
+
+    image_outputs = []
+    for image_path in image_paths:
+      transforms_list = [
+          transforms.Resize(
+              224, interpolation=transforms.InterpolationMode.BICUBIC
+          ),
+          transforms.CenterCrop(224),
+          transforms.ToTensor(),
+      ]
+      if not is_depth_or_thermal:
+        transforms_list.append(
+            transforms.Normalize(
+                mean=(0.48145466, 0.4578275, 0.40821073),
+                std=(0.26862954, 0.26130258, 0.27577711),
+            )
+        )
+      data_transform = transforms.Compose(transforms_list)
+      with open(image_path, "rb") as fopen:
+        if is_depth_or_thermal:
+          image = Image.open(fopen).convert("L")
+        else:
+          image = Image.open(fopen).convert("RGB")
+
+      image = data_transform(image).to(device)
+      image_outputs.append(image)
+    return torch.stack(image_outputs, dim=0)
+
+  def inference(
+      self, data: List[Dict[str, Any]], *args, **kwargs
+  ) -> List[Dict[str, Any]]:
+    """Runs inference using the ImageBind model.
+
+    Args:
+      data: A list of processed data samples, with each sample being a
+        dictionary of modality (key): input (value) pairs.
+      *args: Additional inference args.
+      **kwargs: Additional inference kwargs.
+
+    Returns:
+      A list of model outputs, with each output being a dictionary of
+      modality (key): embedding (value) pairs.
+    """
+    output_list = []
+    with torch.no_grad():
+      for inputs in data:
+        if _VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE in inputs:
+          # Allows inference on both image and video data, which both fall under
+          # ModalityType.VISION.
+          video_inputs = {
+              ModalityType.VISION: inputs[
+                  _VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE
+              ]
+          }
+          video_embeddings = self.model(video_inputs)
+          video_embeddings[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE] = (
+              video_embeddings[ModalityType.VISION]
+          )
+          del video_embeddings[ModalityType.VISION]
+          del inputs[_VIDEO_KEY_TO_AVOID_CONFLICT_WITH_IMAGE]
+        else:
+          video_embeddings = {}
+        embeddings = self.model(inputs)
+        embeddings.update(video_embeddings)
+        output_list.append(embeddings)
+    return output_list
+
+  def postprocess(self, output_list: List[Dict[str, Any]]) -> List[Any]:
+    """Postprocesses model outputs for the task of interest.
+
+    For feature embedding generation, returns the embeddings for each modality
+    for each input.
+    For zero-shot classification, generates classification probabilities
+    between the inputs of a pair of modalities for all possible pairings.
+
+    Args:
+      output_list: A list of model outputs, with each output being a dictionary
+        of modality (key): embedding (value) pairs.
+
+    Returns:
+      A list of postprocessed model outputs for the task of interest, with each
+      output corresponding to an input.
+
+    Raises:
+      ValueError: Fewer than two modalities are provided for zero-shot
+      classification, or the task is not supported.
+    """
+    preds = []
+    if self.task == constants.FEATURE_EMBEDDING_GENERATION:
+      for item in output_list:
+        preds.append({k: v.tolist() for k, v in item.items()})
+    elif self.task == constants.ZERO_SHOT_CLASSIFICATION:
+      for item in output_list:
+        modalities = list(item.keys())
+        if len(modalities) < 2:
+          raise ValueError(
+              "Two or more modalities are needed for task"
+              f" {constants.ZERO_SHOT_CLASSIFICATION}."
+          )
+        pairwise_probs = {}
+        for m1 in modalities:
+          for m2 in modalities:
+            if m1 == m2:
+              continue
+            probs = torch.softmax(item[m1] @ item[m2].T, dim=-1)
+            pairwise_probs[
+                f"Classify each input in {m1} (row) against inputs in"
+                f" {m2} (column)"
+            ] = probs.tolist()
+        preds.append(pairwise_probs)
+    else:
+      raise ValueError(f"Task {self.task} is not supported by the handler.")
+    return preds
@@ -0,0 +1,151 @@
+# This Dockerfile converts JAX vision transformer model to
+# tensorflow saved model format.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="jax-vit-model-conversion:${USER}-test"
+# docker build -f model_oss/jax_vision_transformer/dockerfile/jax_vit_model_conversion.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+
+FROM tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# 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
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Get 'vision_transformer' repository from github.
+RUN git clone https://github.com/google-research/vision_transformer
+# Set current directory to the downloaded 'vision_transformer' repository.
+WORKDIR ./vision_transformer
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard e66b4732d44504251197a3da3f5949f3f3ce9ca6
+
+# Install required libs
+RUN pip install --upgrade pip
+# The following pip installs are pinned down versions of those inside
+# vit_jax/requirements.txt file.
+# NOTE: Using `no-deps` flag to avoid overwriting of
+# dependent library versions. For example,
+# both `chex` and `jax` can overwrite each others
+# `jax-lib` version.
+RUN pip install --no-deps absl-py==1.4.0
+RUN pip install --no-deps aqtp==0.0.10
+RUN pip install --no-deps array-record==0.2.0
+RUN pip install --no-deps astunparse==1.6.3
+RUN pip install --no-deps cached-property==1.5.2
+RUN pip install --no-deps cachetools==5.3.0
+RUN pip install --no-deps certifi==2019.11.28
+RUN pip install --no-deps chardet==3.0.4
+RUN pip install --no-deps chex==0.1.7
+RUN pip install --no-deps click==8.1.3
+RUN pip install --no-deps cloudpickle==2.2.1
+RUN pip install --no-deps clu==0.0.9
+RUN pip install --no-deps contextlib2==21.6.0
+RUN pip install --no-deps dacite==1.8.1
+RUN pip install --no-deps dbus-python==1.2.16
+RUN pip install --no-deps decorator==5.1.1
+RUN pip install --no-deps dm-tree==0.1.8
+RUN pip install --no-deps einops==0.6.1
+RUN pip install --no-deps etils==1.3.0
+RUN pip install --no-deps flatbuffers==23.3.3
+RUN pip install --no-deps flax==0.6.10
+RUN pip install --no-deps git+https://github.com/google/flaxformer@9adaa4467cf17703949b9f537c3566b99de1b416
+RUN pip install --no-deps gast==0.4.0
+RUN pip install --no-deps google-auth==2.16.2
+RUN pip install --no-deps google-auth-oauthlib==0.4.6
+RUN pip install --no-deps google-pasta==0.2.0
+RUN pip install --no-deps googleapis-common-protos==1.59.0
+RUN pip install --no-deps grpcio==1.51.3
+RUN pip install --no-deps h5py==3.8.0
+RUN pip install --no-deps idna==2.8
+RUN pip install --no-deps importlib-metadata==6.1.0
+RUN pip install --no-deps importlib-resources==5.12.0
+RUN pip install --no-deps keras==2.12.0
+RUN pip install --no-deps libclang==16.0.0
+RUN pip install --no-deps Markdown==3.4.3
+RUN pip install --no-deps markdown-it-py==2.2.0
+RUN pip install --no-deps MarkupSafe==2.1.2
+RUN pip install --no-deps mdurl==0.1.2
+RUN pip install --no-deps ml-collections==0.1.1
+RUN pip install --no-deps msgpack==1.0.5
+RUN pip install --no-deps nest-asyncio==1.5.6
+RUN pip install --no-deps numpy==1.23.5
+RUN pip install --no-deps oauthlib==3.2.2
+RUN pip install --no-deps opt-einsum==3.3.0
+RUN pip install --no-deps optax==0.1.5
+RUN pip install --no-deps orbax-checkpoint==0.1.6
+RUN pip install --no-deps packaging==23.0
+RUN pip install --no-deps pandas==2.0.1
+RUN pip install --no-deps pip==23.1.2
+RUN pip install --no-deps promise==2.3
+RUN pip install --no-deps protobuf==4.22.1
+RUN pip install --no-deps psutil==5.9.5
+RUN pip install --no-deps pyasn1==0.4.8
+RUN pip install --no-deps pyasn1-modules==0.2.8
+RUN pip install --no-deps Pygments==2.15.1
+RUN pip install --no-deps PyGObject==3.36.0
+RUN pip install --no-deps python-apt==2.0.1+ubuntu0.20.4.1
+RUN pip install --no-deps python-dateutil==2.8.2
+RUN pip install --no-deps pytz==2023.3
+RUN pip install --no-deps PyYAML==6.0
+RUN pip install --no-deps requests==2.22.0
+RUN pip install --no-deps requests-oauthlib==1.3.1
+RUN pip install --no-deps requests-unixsocket==0.2.0
+RUN pip install --no-deps rich==13.3.5
+RUN pip install --no-deps rsa==4.9
+RUN pip install --no-deps scipy==1.10.1
+RUN pip install --no-deps setuptools==67.6.0
+RUN pip install --no-deps six==1.14.0
+RUN pip install --no-deps tensorboard==2.12.0
+RUN pip install --no-deps tensorboard-data-server==0.7.0
+RUN pip install --no-deps tensorboard-plugin-wit==1.8.1
+RUN pip install --no-deps tensorflow==2.12.0
+RUN pip install --no-deps tensorflow-cpu==2.12.0
+RUN pip install --no-deps tensorflow-datasets==4.9.2
+RUN pip install --no-deps tensorflow-estimator==2.12.0
+RUN pip install --no-deps tensorflow-hub==0.13.0
+RUN pip install --no-deps tensorflow-io-gcs-filesystem==0.31.0
+RUN pip install --no-deps tensorflow-metadata==1.13.1
+RUN pip install --no-deps tensorflow-probability==0.20.0
+RUN pip install --no-deps tensorflow-text==2.12.1
+RUN pip install --no-deps tensorstore==0.1.36
+RUN pip install --no-deps termcolor==2.2.0
+RUN pip install --no-deps toml==0.10.2
+RUN pip install --no-deps toolz==0.12.0
+RUN pip install --no-deps tqdm==4.65.0
+RUN pip install --no-deps typing_extensions==4.5.0
+RUN pip install --no-deps tzdata==2023.3
+RUN pip install --no-deps urllib3==1.25.8
+RUN pip install --no-deps Werkzeug==2.2.3
+RUN pip install --no-deps wheel==0.40.0
+RUN pip install --no-deps wrapt==1.14.1
+RUN pip install --no-deps zipp==3.15.0
+# Installing jax at the very end with GPU support.
+# NOTE: Not using `no-deps` flag here because
+# we need CUDA support.
+RUN pip install jax[cuda11_cudnn82]==0.4.6 \
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+ENV PYTHONPATH ./vit_jax
+
+COPY ./model_oss/jax_vision_transformer/vit_jax2tf.py ./
+COPY ./model_oss/jax_vision_transformer/vit_config_without_data.py vit_jax/configs/vit.py
+
+ENTRYPOINT ["python", "vit_jax2tf.py"]
@@ -0,0 +1,149 @@
+# This Dockerfile runs the JAX based Vision transformer training on GPU.
+# See https://github.com/google-research/vision_transformer#running-on-cloud
+# for more details.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="trainn_vit_gpu:${USER}-test"
+# docker build -f model_oss/jax_vision_transformer/dockerfile/train_vit_gpu.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+FROM tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# 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
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Get 'vision_transformer' repository from github.
+RUN git clone https://github.com/google-research/vision_transformer
+# Ser current directory to the downloaded 'vision_transformer' repository.
+WORKDIR ./vision_transformer
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard e66b4732d44504251197a3da3f5949f3f3ce9ca6
+
+# Install required libs
+RUN pip install --upgrade pip
+# The following pip installs are pinned down versions of those inside
+# vit_jax/requirements.txt file.
+# NOTE: Using `no-deps` flag to avoid overwriting of
+# dependent library versions. For example,
+# both `chex` and `jax` can overwrite each others
+# `jax-lib` version.
+RUN pip install --no-deps absl-py==1.4.0
+RUN pip install --no-deps aqtp==0.0.10
+RUN pip install --no-deps array-record==0.2.0
+RUN pip install --no-deps astunparse==1.6.3
+RUN pip install --no-deps cached-property==1.5.2
+RUN pip install --no-deps cachetools==5.3.0
+RUN pip install --no-deps certifi==2019.11.28
+RUN pip install --no-deps chardet==3.0.4
+RUN pip install --no-deps chex==0.1.7
+RUN pip install --no-deps click==8.1.3
+RUN pip install --no-deps cloudpickle==2.2.1
+RUN pip install --no-deps clu==0.0.9
+RUN pip install --no-deps contextlib2==21.6.0
+RUN pip install --no-deps dacite==1.8.1
+RUN pip install --no-deps dbus-python==1.2.16
+RUN pip install --no-deps decorator==5.1.1
+RUN pip install --no-deps dm-tree==0.1.8
+RUN pip install --no-deps einops==0.6.1
+RUN pip install --no-deps etils==1.3.0
+RUN pip install --no-deps flatbuffers==23.3.3
+RUN pip install --no-deps flax==0.6.10
+RUN pip install --no-deps git+https://github.com/google/flaxformer@9adaa4467cf17703949b9f537c3566b99de1b416
+RUN pip install --no-deps gast==0.4.0
+RUN pip install --no-deps google-auth==2.16.2
+RUN pip install --no-deps google-auth-oauthlib==0.4.6
+RUN pip install --no-deps google-pasta==0.2.0
+RUN pip install --no-deps googleapis-common-protos==1.59.0
+RUN pip install --no-deps grpcio==1.51.3
+RUN pip install --no-deps h5py==3.8.0
+RUN pip install --no-deps idna==2.8
+RUN pip install --no-deps importlib-metadata==6.1.0
+RUN pip install --no-deps importlib-resources==5.12.0
+RUN pip install --no-deps keras==2.12.0
+RUN pip install --no-deps libclang==16.0.0
+RUN pip install --no-deps Markdown==3.4.3
+RUN pip install --no-deps markdown-it-py==2.2.0
+RUN pip install --no-deps MarkupSafe==2.1.2
+RUN pip install --no-deps mdurl==0.1.2
+RUN pip install --no-deps ml-collections==0.1.1
+RUN pip install --no-deps msgpack==1.0.5
+RUN pip install --no-deps nest-asyncio==1.5.6
+RUN pip install --no-deps numpy==1.23.5
+RUN pip install --no-deps oauthlib==3.2.2
+RUN pip install --no-deps opt-einsum==3.3.0
+RUN pip install --no-deps optax==0.1.5
+RUN pip install --no-deps orbax-checkpoint==0.1.6
+RUN pip install --no-deps packaging==23.0
+RUN pip install --no-deps pandas==2.0.1
+RUN pip install --no-deps pip==23.1.2
+RUN pip install --no-deps promise==2.3
+RUN pip install --no-deps protobuf==4.22.1
+RUN pip install --no-deps psutil==5.9.5
+RUN pip install --no-deps pyasn1==0.4.8
+RUN pip install --no-deps pyasn1-modules==0.2.8
+RUN pip install --no-deps Pygments==2.15.1
+RUN pip install --no-deps PyGObject==3.36.0
+RUN pip install --no-deps python-apt==2.0.1+ubuntu0.20.4.1
+RUN pip install --no-deps python-dateutil==2.8.2
+RUN pip install --no-deps pytz==2023.3
+RUN pip install --no-deps PyYAML==6.0
+RUN pip install --no-deps requests==2.22.0
+RUN pip install --no-deps requests-oauthlib==1.3.1
+RUN pip install --no-deps requests-unixsocket==0.2.0
+RUN pip install --no-deps rich==13.3.5
+RUN pip install --no-deps rsa==4.9
+RUN pip install --no-deps scipy==1.10.1
+RUN pip install --no-deps setuptools==67.6.0
+RUN pip install --no-deps six==1.14.0
+RUN pip install --no-deps tensorboard==2.12.0
+RUN pip install --no-deps tensorboard-data-server==0.7.0
+RUN pip install --no-deps tensorboard-plugin-wit==1.8.1
+RUN pip install --no-deps tensorflow==2.12.0
+RUN pip install --no-deps tensorflow-cpu==2.12.0
+RUN pip install --no-deps tensorflow-datasets==4.9.2
+RUN pip install --no-deps tensorflow-estimator==2.12.0
+RUN pip install --no-deps tensorflow-hub==0.13.0
+RUN pip install --no-deps tensorflow-io-gcs-filesystem==0.31.0
+RUN pip install --no-deps tensorflow-metadata==1.13.1
+RUN pip install --no-deps tensorflow-probability==0.20.0
+RUN pip install --no-deps tensorflow-text==2.12.1
+RUN pip install --no-deps tensorstore==0.1.36
+RUN pip install --no-deps termcolor==2.2.0
+RUN pip install --no-deps toml==0.10.2
+RUN pip install --no-deps toolz==0.12.0
+RUN pip install --no-deps tqdm==4.65.0
+RUN pip install --no-deps typing_extensions==4.5.0
+RUN pip install --no-deps tzdata==2023.3
+RUN pip install --no-deps urllib3==1.25.8
+RUN pip install --no-deps Werkzeug==2.2.3
+RUN pip install --no-deps wheel==0.40.0
+RUN pip install --no-deps wrapt==1.14.1
+RUN pip install --no-deps zipp==3.15.0
+# Installing jax at the very end with GPU support.
+# NOTE: Not using `no-deps` flag here because
+# we need CUDA support.
+RUN pip install jax[cuda11_cudnn82]==0.4.6 \
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+COPY ./model_oss/jax_vision_transformer/vit_config_without_data.py vit_jax/configs/vit.py
+
+ENV PYTHONPATH ./vit_jax
+ENTRYPOINT ["python", "-m", "vit_jax.main"]
@@ -0,0 +1,27 @@
+"""Returns a config for a Vision Transformer model without asking for data."""
+import ml_collections
+from vit_jax.configs import common
+from vit_jax.configs import models
+
+
+def get_config(model: str) -> ml_collections.ConfigDict:
+  """Returns default parameters for finetuning ViT `model`."""
+  config = common.get_config()
+
+  get_model_config = getattr(models, f'get_{model}_config')
+  config.model = get_model_config()
+
+  # These values are often overridden on the command line.
+  config.base_lr = 0.03
+  config.total_steps = 500
+  config.warmup_steps = 100
+  config.pp = ml_collections.ConfigDict()
+  config.pp.train = 'train'
+  config.pp.test = 'test'
+  config.pp.resize = 448
+  config.pp.crop = 384
+
+  # This value MUST be overridden on the command line.
+  config.dataset = ''
+
+  return config
@@ -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"]
@@ -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"]
@@ -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)
@@ -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)
@@ -0,0 +1,71 @@
+FROM pytorch/torchserve:0.9.0-gpu
+
+USER root
+
+# Install tools.
+RUN apt-get update -y --allow-releaseinfo-change && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        git
+
+# 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="llava_serving"
+ENV PATH="/home/model-server/:${PATH}"
+ENV PATH="/usr/local/cuda-12.1/bin:${PATH}"
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
+ENV NVIDIA_VISIBLE_DEVICES=all
+
+# Get 'LLaVA' repository from github.
+RUN git clone https://github.com/haotian-liu/LLaVA /home/model-server/LLaVA
+WORKDIR /home/model-server/LLaVA
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard 7775b12d6b20cd69089be7a18ea02615a59621cd
+
+# Install the package.
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.13.0
+RUN pip install absl-py==2.0.0
+RUN pip install -e .
+
+# Copy model artifacts.
+COPY model_oss/llava/handler.py /home/model-server/handler.py
+COPY model_oss/llava/model_handler_setup.py /home/model-server/model_handler_setup.py
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server
+WORKDIR /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}\n"  \
+    "default_workers_per_model=DEFAULT_WORKERS_PER_MODEL" >> /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.
+# Use $NUM_GPU workers unless overriden by $TS_NUM_WORKERS
+CMD ["TOTAL=$(nvidia-smi", "--list-gpus","|","wc","-l)","&&", "TS_NUM_WORKERS=${TS_NUM_WORKERS:-$TOTAL}","&&", "sed","-i","\"s/DEFAULT_WORKERS_PER_MODEL/$TS_NUM_WORKERS/g\"","/home/model-server/config.properties", "&&", \
+    "torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${MODEL_NAME}=${MODEL_NAME}.mar", \
+     "--model-store", "/home/model-server/model-store"]
@@ -0,0 +1,174 @@
+"""Customer handler for LLava 1.5 OSS model.
+
+The code is based on here: https://github.com/haotian-liu/LLaVA
+handler based on:
+https://github.com/haotian-liu/LLaVA/blob/main/llava/eval/run_llava.py
+There are two supported variant:
+1. liuhaotian/llava-v1.5-13b: 13B params
+2. liuhaotian/llava-v1.5-7b: 7B params
+"""
+
+import os
+import re
+from typing import Any, Dict, List
+
+from llava import constants as llava_constants
+from llava import conversation
+from llava import mm_utils
+from llava.model import builder
+import model_handler_setup
+import torch
+from ts.torch_handler import base_handler
+
+from util import constants
+from util import image_format_converter
+
+
+DEFAULT_MODEL_ID = "liuhaotian/llava-v1.5-7b"
+
+
+class LlavaHandler(base_handler.BaseHandler):
+  """Custom handler for LLava model."""
+
+  def initialize(self, context: Any):
+    """Initializes model, tokenizer, and other components."""
+    self.map_location = model_handler_setup.get_map_location(context=context)
+    self.device = model_handler_setup.get_model_device(
+        map_location=self.map_location, context=context
+    )
+    self.manifest = context.manifest
+    self.model_id = model_handler_setup.get_model_id(
+        default_model_id=DEFAULT_MODEL_ID
+    )
+
+    # Allows 4bit and 8bit quantiziation using BnB nf4.
+    precision = os.environ.get("PRECISION_MODE")
+    load_8bit = precision == constants.PRECISION_MODE_8
+    load_4bit = precision == constants.PRECISION_MODE_4
+
+    self.tokenizer, self.model, self.image_processor, self.context_len = (
+        builder.load_pretrained_model(
+            model_path=self.model_id,
+            model_base=None,
+            model_name=mm_utils.get_model_name_from_path(self.model_id),
+            load_8bit=load_8bit,
+            load_4bit=load_4bit,
+        )
+    )
+
+  def preprocess(self, data: List[Dict[str, Any]]) -> Any:
+    """Runs the preprocessing to tokenize image and the prompt."""
+    if len(data) > 1:
+      raise ValueError(
+          "LLava original repo currently does not support batch inference."
+          " https://github.com/haotian-liu/LLaVA/issues/754"
+      )
+    data = data[0]
+    prompt, base64_image = data["prompt"], data["base64_image"]
+
+    # Adds proper image token to the prompt.
+    image_token_se = (
+        llava_constants.DEFAULT_IM_START_TOKEN
+        + llava_constants.DEFAULT_IMAGE_TOKEN
+        + llava_constants.DEFAULT_IM_END_TOKEN
+    )
+    if llava_constants.IMAGE_PLACEHOLDER in prompt:
+      if self.model.config.mm_use_im_start_end:
+        prompt = re.sub(
+            llava_constants.IMAGE_PLACEHOLDER, image_token_se, prompt
+        )
+      else:
+        prompt = re.sub(
+            llava_constants.IMAGE_PLACEHOLDER,
+            llava_constants.DEFAULT_IMAGE_TOKEN,
+            prompt,
+        )
+    else:
+      if self.model.config.mm_use_im_start_end:
+        prompt = image_token_se + "\n" + prompt
+      else:
+        prompt = llava_constants.DEFAULT_IMAGE_TOKEN + "\n" + prompt
+
+    # Formats the prompt as a conversation to be fed to the model.
+    conv = conversation.conv_llava_v1.copy()
+    conv.append_message(role=conv.roles[0], message=prompt)
+    conv.append_message(role=conv.roles[1], message=None)
+    prompt = conv.get_prompt()
+
+    # Tokenizes the prompt that includes special image token as well.
+    input_ids = (
+        mm_utils.tokenizer_image_token(
+            prompt=prompt,
+            tokenizer=self.tokenizer,
+            image_token_index=llava_constants.IMAGE_TOKEN_INDEX,
+            return_tensors="pt",
+        )
+        .unsqueeze(0)
+        .to(self.device)
+    )
+
+    images = [
+        image_format_converter.base64_to_image(image_str=base64_image).convert(
+            "RGB"
+        )
+    ]
+    # Gets the image embedding.
+    images_tensor = mm_utils.process_images(
+        images=images,
+        image_processor=self.image_processor,
+        model_cfg=self.model.config,
+    ).to(self.device, dtype=torch.float16)
+
+    self.stop_str = conversation.conv_llava_v1.sep2
+    self.keywords = [self.stop_str]
+
+    return input_ids, images_tensor
+
+  def inference(
+      self, input_ids: List[torch.Tensor], images_tensor: torch.Tensor
+  ) -> List[torch.Tensor]:
+    """Runs the inference."""
+    stopping_criteria = mm_utils.KeywordsStoppingCriteria(
+        keywords=self.keywords, tokenizer=self.tokenizer, input_ids=input_ids
+    )
+
+    with torch.inference_mode():
+      output_ids = self.model.generate(
+          input_ids=input_ids,
+          images=images_tensor,
+          do_sample=False,
+          temperature=0,
+          top_p=None,
+          num_beams=1,
+          max_new_tokens=512,
+          use_cache=True,
+          stopping_criteria=[stopping_criteria],
+      )
+
+    return output_ids
+
+  def postprocess(
+      self, output_ids: List[torch.Tensor], input_token_len: int
+  ) -> List[str]:
+    """Runs the postprocessing to convert token ids to string."""
+    outputs = self.tokenizer.batch_decode(
+        output_ids[:, input_token_len:], skip_special_tokens=True
+    )[0]
+    outputs = outputs.strip()
+    if outputs.endswith(self.stop_str):
+      outputs = outputs[: -len(self.stop_str)]
+    outputs = outputs.strip()
+
+    return [outputs]
+
+  def handle(self, data: List[Dict[str, Any]], context: Any) -> List[str]:
+    """Handles an incoming request by passing it through `preprocess`, `inference`, and `postprocess`."""
+    input_ids, images_tensor = self.preprocess(data=data)
+    model_output = self.inference(
+        input_ids=input_ids, images_tensor=images_tensor
+    )
+
+    input_token_len = input_ids.shape[1]
+    return self.postprocess(
+        output_ids=model_output, input_token_len=input_token_len
+    )
@@ -0,0 +1,77 @@
+"""Common utility functions for setting up and initializing the model and the handler."""
+
+import logging
+import os
+from typing import Any
+
+import torch
+
+from util import constants
+from util import fileutils
+
+
+def get_model_id(default_model_id: str) -> str:
+  """Gets a model id or a local model path.
+
+  Args:
+    default_model_id: Default model id for the corresponding model set in the
+      handler.
+
+  Returns:
+    str: model id or a local model path.
+  """
+  # The model id can be either:
+  # 1) a huggingface model card id, like "Salesforce/blip", or
+  # 2) a GCS path to the model files, like "gs://foo/bar".
+  # If it's a model card id, the model will be loaded from huggingface.
+  model_id = (
+      default_model_id
+      if os.environ.get("MODEL_ID") is None
+      else os.environ["MODEL_ID"]
+  )
+
+  # Else it will be downloaded from GCS to local first.
+  # Since the transformers from_pretrained API can't read from GCS.
+  if model_id.startswith(constants.GCS_URI_PREFIX):
+    gcs_path = model_id[len(constants.GCS_URI_PREFIX) :]
+    local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
+    logging.info("Download %s to %s", model_id, local_model_dir)
+    fileutils.download_gcs_dir_to_local(model_id, local_model_dir)
+    model_id = local_model_dir
+
+  return model_id
+
+
+def get_map_location(context: Any) -> str:
+  """Gets model map location.
+
+  Args:
+    context: Torchserve worker context.
+
+  Returns:
+    str: Mapping location.
+  """
+  properties = context.system_properties
+  return (
+      "cuda"
+      if torch.cuda.is_available() and properties.get("gpu_id") is not None
+      else "cpu"
+  )
+
+
+def get_model_device(map_location: str, context: Any) -> torch.device:
+  """Gets model accelerator device.
+
+  Args:
+    map_location: Model map location.
+    context: TorchServe worker context.
+
+  Returns:
+    torch.Device: Device to load the model into.
+  """
+  properties = context.system_properties
+  return torch.device(
+      map_location + ":" + str(properties.get("gpu_id"))
+      if torch.cuda.is_available() and properties.get("gpu_id") is not None
+      else map_location
+  )
@@ -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]"
@@ -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
@@ -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"]
@@ -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"]
@@ -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"]
@@ -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)
@@ -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
+  ]
@@ -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)
@@ -0,0 +1,564 @@
+"""Common util functions for notebook."""
+
+import base64
+import datetime
+import io
+import json
+import os
+import subprocess
+from typing import Any, Dict, Sequence
+
+from google.cloud import storage
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+import requests
+import tensorflow as tf
+import yaml
+
+
+GCS_URI_PREFIX = "gs://"
+CHECKPOINT_BUCKET = "gs://model_garden_checkpoints"
+
+
+def convert_numpy_array_to_byte_string_via_tf_tensor(
+    np_array: np.ndarray,
+) -> str:
+  """Serializes a numpy array to tensor bytes.
+
+  Args:
+    np_array: A numpy array.
+
+  Returns:
+    A tensor bytes.
+  """
+  tensor_array = tf.convert_to_tensor(np_array)
+  tensor_byte_string = tf.io.serialize_tensor(tensor_array)
+  return tensor_byte_string.numpy()
+
+
+def get_jpeg_bytes(local_image_path: str, new_width: int = -1) -> bytes:
+  """Returns jpeg bytes given an image path and resizes if required.
+
+  Args:
+    local_image_path: A string of local image path.
+    new_width: An integer of new image width.
+
+  Returns:
+    A jpeg bytes.
+  """
+  image = Image.open(local_image_path)
+  if new_width <= 0:
+    new_image = image
+  else:
+    width, height = image.size
+    print("original input image size: ", width, " , ", height)
+    new_height = int(height * new_width / width)
+    print("new input image size: ", new_width, " , ", new_height)
+    new_image = image.resize((new_width, new_height))
+  buffered = io.BytesIO()
+  new_image.save(buffered, format="JPEG")
+  return buffered.getvalue()
+
+
+def gcs_fuse_path(path: str) -> str:
+  """Try to convert path to gcsfuse path if it starts with gs:// else do not modify it.
+
+  Args:
+    path: A string of path.
+
+  Returns:
+    A gcsfuse path.
+  """
+  path = path.strip()
+  if path.startswith("gs://"):
+    return "/gcs/" + path[5:]
+  return path
+
+
+def get_job_name_with_datetime(prefix: str) -> str:
+  """Gets a job name by adding current time to prefix.
+
+  Args:
+    prefix: A string of job name prefix.
+
+  Returns:
+    A job name.
+  """
+  now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+  job_name = f"{prefix}-{now}".replace("_", "-")
+  return job_name
+
+
+def create_job_name(prefix: str) -> str:
+  """Creates a job name.
+
+  Args:
+    prefix: A string of job name prefix.
+
+  Returns:
+    A job name.
+  """
+  user = os.environ.get("USER")
+  now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+  job_name = f"{prefix}-{user}-{now}".replace("_", "-")
+  return job_name
+
+
+def save_subset_annotation(
+    input_annotation_path: str, output_annotation_path: str
+):
+  """Saves a subset of COCO annotation json file with CCA 4.0 license.
+
+  Args:
+    input_annotation_path: A string of input annotation path.
+    output_annotation_path: A string of output annotation path.
+  """
+
+  with open(input_annotation_path) as f:
+    coco_json = json.load(f)
+
+  img_ids = set()
+  images = []
+  annotations = []
+
+  for img in coco_json["images"]:
+    if img["license"] in [4, 5]:  # CCA 4.0 license.
+      img_ids.add(img["id"])
+      images.append(img)
+
+  for ann in coco_json["annotations"]:
+    if ann["image_id"] in img_ids:
+      annotations.append(ann)
+
+  new_json = {
+      "info": coco_json["info"],
+      "licenses": coco_json["licenses"],
+      "images": images,
+      "annotations": annotations,
+      "categories": coco_json["categories"],
+  }
+
+  with open(output_annotation_path, "w") as f:
+    json.dump(new_json, f)
+
+
+def image_to_base64(image: Any, image_format: str = "JPEG") -> str:
+  """Converts an image to base64.
+
+  Args:
+    image: A PIL.Image instance.
+    image_format: A string of image format.
+
+  Returns:
+    A base64 string.
+  """
+  buffer = io.BytesIO()
+  image.save(buffer, format=image_format)
+  image_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
+  return image_str
+
+
+def base64_to_image(image_str: str) -> Any:
+  """Convert base64 encoded string to an image.
+
+  Args:
+    image_str: A string of base64 encoded image.
+
+  Returns:
+    A PIL.Image instance.
+  """
+  image = Image.open(io.BytesIO(base64.b64decode(image_str)))
+  return image
+
+
+def image_grid(imgs: Sequence[Any], rows: int = 2, cols: int = 2) -> Any:
+  """Creates an image grid.
+
+  Args:
+    imgs: A list of PIL.Image instances.
+    rows: An integer of number of rows.
+    cols: An integer of number of columns.
+
+  Returns:
+    A PIL.Image instance.
+  """
+  w, h = imgs[0].size
+  grid = Image.new(
+      mode="RGB", size=(cols * w + 10 * cols, rows * h), color=(255, 255, 255)
+  )
+  for i, img in enumerate(imgs):
+    grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))
+  return grid
+
+
+def display_image(image: Any):
+  """Displays an image.
+
+  Args:
+    image: A PIL.Image instance.
+  """
+  _ = plt.figure(figsize=(20, 15))
+  plt.grid(False)
+  plt.imshow(image)
+
+
+def download_gcs_file_to_local(gcs_uri: str, local_path: str):
+  """Download a gcs file to a local path.
+
+  Args:
+    gcs_uri: A string of file path on GCS.
+    local_path: A string of local file path.
+  """
+  if not gcs_uri.startswith(GCS_URI_PREFIX):
+    raise ValueError(
+        f"{gcs_uri} is not a GCS path starting with {GCS_URI_PREFIX}."
+    )
+  client = storage.Client()
+  os.makedirs(os.path.dirname(local_path), exist_ok=True)
+  with open(local_path, "wb") as f:
+    client.download_blob_to_file(gcs_uri, f)
+
+
+def download_image(url: str) -> str:
+  """Downloads an image from the given URL.
+
+  Args:
+    url: The URL of the image to download.
+
+  Returns:
+    base64 encoded image.
+  """
+  response = requests.get(url)
+  return Image.open(io.BytesIO(response.content))
+
+
+def resize_image(image: Any, new_width: int = 1000) -> Any:
+  """Resizes an image to a certain width.
+
+  Args:
+    image: The image which has to be resized.
+    new_width: New width of the image.
+
+  Returns:
+    New resized image.
+  """
+  width, height = image.size
+  new_height = int(height * new_width / width)
+  new_img = image.resize((new_width, new_height))
+  return new_img
+
+
+def load_img(path: str) -> Any:
+  """Reads image from path and return PIL.Image instance.
+
+  Args:
+    path: A string of image path.
+
+  Returns:
+    A PIL.Image instance.
+  """
+  img = tf.io.read_file(path)
+  img = tf.image.decode_jpeg(img, channels=3)
+  return Image.fromarray(np.uint8(img)).convert("RGB")
+
+
+def decode_image(
+    image_str_tensor: tf.string, new_height: int, new_width: int
+) -> tf.float32:
+  """Converts and resizes image bytes to image tensor.
+
+  Args:
+    image_str_tensor: A string of image bytes.
+    new_height: An integer of new image height.
+    new_width: An integer of new image width.
+
+  Returns:
+    An image tensor.
+  """
+  image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)
+  image = tf.image.resize(image, (new_height, new_width))
+  return image
+
+
+def get_label_map(label_map_yaml_filepath: str) -> Dict[int, str]:
+  """Returns class id to label mapping given a filepath to the label map.
+
+  Args:
+    label_map_yaml_filepath: A string of label map yaml file path.
+
+  Returns:
+    A dictionary of class id to label mapping.
+  """
+  with tf.io.gfile.GFile(label_map_yaml_filepath, "rb") as input_file:
+    label_map = yaml.safe_load(input_file.read())["label_map"]
+  return label_map
+
+
+def get_prediction_instances(test_filepath: str, new_width: int = -1) -> Any:
+  """Generate instance from image path to pass to Vertex AI Endpoint for prediction.
+
+  Args:
+    test_filepath: A string of test image path.
+    new_width: An integer of new image width.
+
+  Returns:
+    A list of instances.
+  """
+  if new_width <= 0:
+    with tf.io.gfile.GFile(test_filepath, "rb") as input_file:
+      encoded_string = base64.b64encode(input_file.read()).decode("utf-8")
+  else:
+    img = load_img(test_filepath)
+    width, height = img.size
+    print("original input image size: ", width, " , ", height)
+    new_height = int(height * new_width / width)
+    new_img = img.resize((new_width, new_height))
+    print("resized input image size: ", new_width, " , ", new_height)
+    buffered = io.BytesIO()
+    new_img.save(buffered, format="JPEG")
+    encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+  instances = [{
+      "encoded_image": {"b64": encoded_string},
+  }]
+  return instances
+
+
+def vqa_predict(
+    endpoint: Any,
+    question_prompts: Sequence[str],
+    image: Any,
+    language_code: str = "en",
+    new_width: int = 1000,
+) -> Sequence[str]:
+  """Predicts the answer to a question about an image using an Endpoint."""
+  # Resize and convert image to base64 string.
+  resized_image = resize_image(image, new_width)
+  resized_image_base64 = image_to_base64(resized_image)
+
+  instances = []
+  if question_prompts:
+    # Format question prompt
+    question_prompt_format = "answer {} {}\n"
+    for question_prompt in question_prompts:
+      if question_prompt:
+        instances.append({
+            "prompt": question_prompt_format.format(
+                language_code, question_prompt
+            ),
+            "image": resized_image_base64,
+        })
+  else:
+    instances.append({
+        "image": resized_image_base64,
+    })
+
+  response = endpoint.predict(instances=instances)
+  return [pred.get("response") for pred in response.predictions]
+
+
+def caption_predict(
+    endpoint: Any,
+    language_code: str,
+    image: Any,
+    caption_prompt: bool = False,
+    new_width: int = 1000,
+) -> str:
+  """Predicts a caption for a given image using an Endpoint."""
+  # Resize and convert image to base64 string.
+  resized_image = resize_image(image, new_width)
+  resized_image_base64 = image_to_base64(resized_image)
+
+  instance = {"image": resized_image_base64}
+
+  if caption_prompt:
+    # Format caption prompt
+    caption_prompt_format = "caption {}\n"
+    instance["prompt"] = caption_prompt_format.format(language_code)
+
+  instances = [instance]
+  response = endpoint.predict(instances=instances)
+  return response.predictions[0].get("response")
+
+
+def ocr_predict(
+    endpoint: Any,
+    ocr_prompt: str,
+    image: Any,
+    new_width: int = 1000,
+) -> str:
+  """Extracts text from a given image using an Endpoint."""
+  # Resize and convert image to base64 string.
+  resized_image = resize_image(image, new_width)
+  resized_image_base64 = image_to_base64(resized_image)
+
+  instance = {"image": resized_image_base64}
+  if ocr_prompt:
+    instance["prompt"] = ocr_prompt
+  instances = [instance]
+
+  response = endpoint.predict(instances=instances)
+  return response.predictions[0].get("response")
+
+
+def detect_predict(
+    endpoint: Any,
+    detect_prompt: str,
+    image: Any,
+    new_width: int = 1000,
+) -> str:
+  """Predicts the answer to a question about an image using an Endpoint."""
+  # Resize and convert image to base64 string.
+  resized_image = resize_image(image, new_width)
+  resized_image_base64 = image_to_base64(resized_image)
+
+  instance = {"image": resized_image_base64}
+  if detect_prompt:
+    instance["prompt"] = detect_prompt
+  instances = [instance]
+
+  response = endpoint.predict(instances=instances)
+  return response.predictions[0].get("response")
+
+
+def get_quota(project_id: str, region: str, resource_id: str) -> int:
+  """Returns the quota for a resource in a region.
+
+  Args:
+    project_id: The project id.
+    region: The region.
+    resource_id: The resource id.
+
+  Returns:
+    The quota for the resource in the region. Returns -1 if can not figure out
+    the quota.
+
+  Raises:
+    RuntimeError: If the command to get quota fails.
+  """
+  service_endpoint = "aiplatform.googleapis.com"
+
+  command = (
+      "gcloud alpha services quota list"
+      f" --service={service_endpoint} --consumer=projects/{project_id}"
+      f" --filter='{service_endpoint}/{resource_id}' --format=json"
+  )
+  process = subprocess.run(
+      command, shell=True, capture_output=True, text=True, check=True
+  )
+  if process.returncode == 0:
+    quota_data = json.loads(process.stdout)
+  else:
+    raise RuntimeError(f"Error fetching quota data: {process.stderr}")
+
+  if not quota_data or "consumerQuotaLimits" not in quota_data[0]:
+    return -1
+  if (
+      not quota_data[0]["consumerQuotaLimits"]
+      or "quotaBuckets" not in quota_data[0]["consumerQuotaLimits"][0]
+  ):
+    return -1
+  all_regions_data = quota_data[0]["consumerQuotaLimits"][0]["quotaBuckets"]
+  for region_data in all_regions_data:
+    if (
+        region_data.get("dimensions")
+        and region_data["dimensions"]["region"] == region
+    ):
+      if "effectiveLimit" in region_data:
+        return int(region_data["effectiveLimit"])
+      else:
+        return 0
+  return -1
+
+
+def get_resource_id(
+    accelerator_type: str,
+    is_for_training: bool,
+    is_restricted_image: bool = False,
+) -> str:
+  """Returns the resource id for a given accelerator type and the use case.
+
+  Args:
+    accelerator_type: The accelerator type.
+    is_for_training: Whether the resource is used for training. Set false for
+      serving use case.
+    is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
+
+  Returns:
+    The resource id.
+  """
+  default_training_accelerator_map = {
+      "NVIDIA_TESLA_V100": "custom_model_training_nvidia_v100_gpus",
+      "NVIDIA_L4": "custom_model_training_nvidia_l4_gpus",
+      "NVIDIA_TESLA_A100": "custom_model_training_nvidia_a100_gpus",
+      "NVIDIA_A100_80GB": "custom_model_training_nvidia_a100_80gb_gpus",
+      "NVIDIA_H100_80GB": "custom_model_training_nvidia_h100_gpus",
+      "NVIDIA_TESLA_T4": "custom_model_training_nvidia_t4_gpus",
+      "TPU_V5e": "custom_model_training_tpu_v5e",
+      "TPU_V3": "custom_model_training_tpu_v3",
+  }
+  restricted_image_training_accelerator_map = {
+      "NVIDIA_A100_80GB": "restricted_image_training_nvidia_a100_80gb_gpus",
+  }
+  serving_accelerator_map = {
+      "NVIDIA_TESLA_V100": "custom_model_serving_nvidia_v100_gpus",
+      "NVIDIA_L4": "custom_model_serving_nvidia_l4_gpus",
+      "NVIDIA_TESLA_A100": "custom_model_serving_nvidia_a100_gpus",
+      "NVIDIA_A100_80GB": "custom_model_serving_nvidia_a100_80gb_gpus",
+      "NVIDIA_H100_80GB": "custom_model_serving_nvidia_h100_gpus",
+      "NVIDIA_TESLA_T4": "custom_model_serving_nvidia_t4_gpus",
+      "TPU_V5e": "custom_model_serving_tpu_v5e",
+  }
+  if is_for_training:
+    training_accelerator_map = (
+        restricted_image_training_accelerator_map
+        if is_restricted_image
+        else default_training_accelerator_map
+    )
+    if accelerator_type in training_accelerator_map:
+      return training_accelerator_map[accelerator_type]
+    else:
+      raise ValueError(
+          f"Could not find accelerator type: {accelerator_type} for training."
+      )
+  else:
+    if accelerator_type in serving_accelerator_map:
+      return serving_accelerator_map[accelerator_type]
+    else:
+      raise ValueError(
+          f"Could not find accelerator type: {accelerator_type} for serving."
+      )
+
+
+def check_quota(
+    project_id: str,
+    region: str,
+    accelerator_type: str,
+    accelerator_count: int,
+    is_for_training: bool,
+    is_restricted_image: bool = False,
+):
+  """Checks if the project and the region has the required quota."""
+  resource_id = get_resource_id(
+      accelerator_type, is_for_training, is_restricted_image
+  )
+  quota = get_quota(project_id, region, resource_id)
+  quota_request_instruction = (
+      "Either use "
+      "a different region or request additional quota. Follow "
+      "instructions here "
+      "https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota"
+      " to check quota in a region or request additional quota for "
+      "your project."
+  )
+  if quota == -1:
+    raise ValueError(
+        f"Quota not found for: {resource_id} in {region}."
+        f" {quota_request_instruction}"
+    )
+  if quota < accelerator_count:
+    raise ValueError(
+        f"Quota not enough for {resource_id} in {region}: {quota} <"
+        f" {accelerator_count}. {quota_request_instruction}"
+    )
@@ -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"]
@@ -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
@@ -0,0 +1,172 @@
+"""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 util import constants
+from util import fileutils
+from util import image_format_converter
+
+
+# Supported checkpoint&model pairs:
+# https://github.com/mlfoundations/open_clip#pretrained-model-interface
+_DEFAULT_MODEL = "RN50"
+_BIOMED_CLIP_MODEL = "microsoft/BiomedCLIP"
+_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."""
+
+  @enum.unique
+  class Precision(enum.Enum):
+    AMP = "amp"
+    AMP_BF16 = "amp_bf16"
+    AMP_BFLOAT16 = "amp_bfloat16"
+    # For the difference between floating points and "pure" floating points, see
+    # https://github.com/mlfoundations/open_clip/blob/0142d279298a4ca0138316286f775fe9d7bdbb94/src/open_clip/factory.py#L232C58-L232C58
+    BF16 = "bf16"
+    FP16 = "fp16"
+    PURE_BF16 = "pure_bf16"
+    PURE_FP16 = "pure_fp16"
+    FP32 = "fp32"
+
+  _DEFAULT_PRECISION = Precision.AMP
+
+  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
+
+    self.model_name = os.environ.get("MODEL", None)
+    if not self.model_name:
+      self.model_name = os.environ.get("MODEL_ID", _DEFAULT_MODEL)
+    precision = os.environ.get("PRECISION", self._DEFAULT_PRECISION)
+    checkpoint = os.environ.get("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,
+        self.model_name,
+        precision,
+        checkpoint,
+    )
+
+    if fileutils.is_gcs_path(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_from_pretrained(
+        self.model_name, pretrained=checkpoint, precision=precision
+    )
+    self.model.to(self.device)
+    self.tokenizer = open_clip.get_tokenizer(self.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 _biomedclip_inference(
+      self, data: List[Dict[str, Any]], *args, **kwargs
+  ) -> List[List[float]]:
+    """Inference for BiomedCLIP model."""
+    texts = torch.stack(
+        [item[_TEXT_KEY][0] for item in data if _TEXT_KEY in item]
+    ).to(self.map_location)
+    images = torch.stack(
+        [item[_IMAGE_KEY][0] for item in data if _IMAGE_KEY in item]
+    ).to(self.map_location)
+    if texts.shape[0] == 0 or images.shape[0] == 0:
+      return []
+    with torch.no_grad():
+      image_features, text_features, logit_scale = self.model(images, texts)
+      logits = (
+          (logit_scale * image_features @ text_features.t())
+          .detach()
+          .softmax(dim=-1)
+      )
+      return logits.cpu().numpy().tolist()
+
+  def inference(
+      self, data: List[Dict[str, Any]], *args, **kwargs
+  ) -> List[Dict[str, Any]]:
+    if _BIOMED_CLIP_MODEL in self.model_name:
+      return self._biomedclip_inference(data)
+    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."""
+    if _BIOMED_CLIP_MODEL in self.model_name:
+      return features
+    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
@@ -0,0 +1,144 @@
+"""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 typing import List
+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,
+    target_modules: List[str] = constants.CAUSAL_LANGUAGE_MODELING_LORA_TARGET_MODULES,
+    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=target_modules,
+      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)
@@ -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\
+    }\
+  }\
+}
@@ -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 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"]
@@ -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"]
@@ -0,0 +1,332 @@
+"""Custom handler for huggingface/peft models."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=logging-fstring-interpolation
+
+import os
+import time
+from typing import Any, List, Tuple
+
+from absl import logging
+from awq import AutoAWQForCausalLM
+from diffusers import DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline
+from peft import PeftModel
+from PIL import Image
+import psutil
+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
+
+if os.path.exists(constants.SHARED_MEM_DIR):
+  logging.info(
+      "SharedMemorySizeMb: %s",
+      psutil.disk_usage(constants.SHARED_MEM_DIR).free / 1e6,
+  )
+
+# 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
+_MAX_TOKENS_DEFAULT = None
+_TEMPERATURE_DEFAULT = 1.0
+_TOP_P_DEFAULT = 1.0
+_TOP_K_DEFAULT = 10
+
+logging.set_verbosity(os.environ.get("LOG_LEVEL", logging.INFO))
+
+
+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", None)
+    self.model_id = self.base_model_id
+    if not self.base_model_id:
+      self.model_id = os.environ.get("MODEL_ID", "")
+    self.quantization = os.environ.get("QUANTIZATION", None)
+    logging.info(f"Load base model id from MODEL_ID:{self.model_id}.")
+    if not self.model_id:
+      self.model_id = os.environ.get("AIP_STORAGE_URI", "")
+      logging.info(f"Load base model id from AIP_STORAGE_URI: {self.model_id}.")
+    if not self.model_id:
+      raise ValueError("Base model id is must be set.")
+    if fileutils.is_gcs_path(self.model_id):
+      fileutils.download_gcs_dir_to_local(
+          self.model_id,
+          constants.LOCAL_BASE_MODEL_DIR,
+          skip_hf_model_bin=True,
+      )
+      self.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.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
+    start_time = time.perf_counter()
+    logging.info("Started PEFT handler initialization at: %s", start_time)
+    if self.task == TEXT_TO_IMAGE_LORA:
+      pipeline = StableDiffusionPipeline.from_pretrained(
+          self.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.")
+      # 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()
+      if self.finetuned_lora_model_path:
+        pipeline.load_lora_weights(self.finetuned_lora_model_path)
+      logging.debug("Initialized the LoRA model for text to image.")
+      self.pipeline = pipeline
+      logging.info("Initialized the text to image pipelines.")
+    elif self.task == SEQUENCE_CLASSIFICATION_LORA:
+      tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+      logging.debug("Initialized the tokenizer for sequence classification.")
+      model = AutoModelForSequenceClassification.from_pretrained(
+          self.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.model_id)
+      logging.debug("Initialized the tokenizer.")
+      if self.task == CAUSAL_LANGUAGE_MODELING_LORA:
+        if self.quantization == constants.AWQ:
+          model = AutoAWQForCausalLM.from_quantized(self.model_id)
+        elif self.quantization == constants.GPTQ or not self.quantization:
+          if self.precision_mode == constants.PRECISION_MODE_32:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                torch_dtype=torch.float32,
+                device_map="auto",
+            )
+          elif self.precision_mode == constants.PRECISION_MODE_16B:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+            )
+          elif self.precision_mode == constants.PRECISION_MODE_16:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                torch_dtype=torch.float16,
+                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.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.model_id,
+                return_dict=True,
+                device_map="auto",
+                torch_dtype=torch.bfloat16,
+                quantization_config=quantization_config,
+            )
+        else:
+          raise ValueError(f"Invalid QUANTIZATION value: {self.quantization}")
+      else:
+        try:
+          model = AutoModelForCausalLM.from_pretrained(
+              self.model_id,
+              torch_dtype=torch.bfloat16,
+              trust_remote_code=True,
+              device_map="auto",
+          )
+        except:  # pylint: disable=bare-except
+          model = AutoModelForCausalLM.from_pretrained(
+              self.model_id,
+              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
+    end_time = time.perf_counter()
+    logging.info("The PEFT handler was initialize at: %s", end_time)
+    logging.info("Handler initiation took %s seconds", end_time - start_time)
+
+  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.
+    # For generation length: `max_length` defines the maximum length of the
+    # sequence to be generated, including both input and output tokens.
+    # `max_length` is overridden by `max_new_tokens` if also set.
+    # `max_new_tokens` defines the maximum number of new tokens to generate,
+    # ignoring the current number of tokens.
+    # Reference:
+    # https://github.com/huggingface/transformers/blob/574a5384557b1aaf98ddb13ea9eb0a0ee8ff2cb2/src/transformers/generation/configuration_utils.py#L69-L73
+    max_length = _MAX_LENGTH_DEFAULT
+    max_tokens = _MAX_TOKENS_DEFAULT
+    temperature = _TEMPERATURE_DEFAULT
+    top_p = _TOP_P_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 "max_tokens" in data[0]:
+      max_tokens = data[0]["max_tokens"]
+    if "temperature" in data[0]:
+      temperature = data[0]["temperature"]
+    if "top_p" in data[0]:
+      top_p = data[0]["top_p"]
+    if "top_k" in data[0]:
+      top_k = data[0]["top_k"]
+
+    return prompts, max_length, max_tokens, temperature, top_p, top_k
+
+  def inference(
+      self, data: Any, *args, **kwargs
+  ) -> Tuple[List[str], List[Image.Image]]:
+    """Runs the inference."""
+    prompts, max_length, max_tokens, temperature, top_p, top_k = data
+    logging.debug(
+        f"Inference prompts={prompts}, max_length={max_length},"
+        f" max_tokens={max_tokens}, temperature={temperature}, top_p={top_p},"
+        f" 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", padding=True)
+      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,
+          max_new_tokens=max_tokens,
+          do_sample=True,
+          temperature=temperature,
+          top_p=top_p,
+          top_k=top_k,
+          num_return_sequences=1,
+          eos_token_id=self.tokenizer.eos_token_id,
+          return_full_text=False,
+      )
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+    return prompts, predicted_results
+
+  def postprocess(self, data: Any) -> List[str]:
+    """Postprocesses output data."""
+    prompts, predicted_results = 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 predicted_results
+      ]
+    elif self.task == SEQUENCE_CLASSIFICATION_LORA:
+      outputs = predicted_results
+    else:
+      outputs = []
+      for prompt, predicted_result in zip(prompts, predicted_results):
+        formatted_output = self._format_text_generation_output(
+            prompt=prompt, output=predicted_result[0]["generated_text"]
+        )
+        outputs.append(formatted_output)
+    return outputs
+
+  def _format_text_generation_output(self, prompt: str, output: str) -> str:
+    """Formats text generation output."""
+    output = output.strip("\n")
+    return f"Prompt:\n{prompt.strip()}\nOutput:\n{output}"
+
+
+# pylint: enable=logging-fstring-interpolation
--- a/Show More
+++ b/Show More