### Setup Big Vision Repository and Install Dependencies

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/uvim_color_task.ipynb

This snippet fetches the Big Vision repository, copies the relevant 'big_vision' directory, and installs the necessary Python dependencies using pip. It ensures the project code is available in the import path.

```shell
# Fetch big_vision repository and move it into the current workdir (import path).
!git clone --depth=1 https://github.com/google-research/big_vision big_vision_repo
!cp -R big_vision_repo/big_vision big_vision
!pip install -qr big_vision/requirements.txt
```

--------------------------------

### Setup Python Virtual Environment and Install Dependencies

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/README.md

Creates a Python virtual environment named 'bv', activates it, navigates to the 'big_vision' directory, upgrades pip, and installs project dependencies from 'requirements.txt'. It also installs the JAX library with TPU support from a specific URL.

```bash
virtualenv bv
source bv/bin/activate
cd big_vision/
pip3 install --upgrade pip
pip3 install -r big_vision/requirements.txt
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```

--------------------------------

### Install big_vision Requirements

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

Installs the necessary Python packages required by the big_vision library. This command reads the requirements from the specified file and installs them using pip.

```bash
!pip install -qr big_vision/big_vision/requirements.txt
```

--------------------------------

### Invoke Stage II Panoptic Training (UViM Pretrained)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/README.md

Starts the stage II panoptic training for the UViM model, utilizing a pretrained checkpoint. It points to the specific configuration file and defines a dynamic working directory. The ':singlehost' flag ensures compatibility with a single-host setup.

```python
python3 -m big_vision.trainers.proj.uvim.train --config big_vision/configs/proj/uvim/train_coco_panoptic_pretrained.py:singlehost --workdir workdirs/`date '+%m-%d_%H%M'`
```

--------------------------------

### Environment Setup for Big Vision Project (Python)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP_demo.ipynb

Sets up the Colab environment by installing the JAX library, cloning the Big Vision repository, installing dependencies, and updating the working directory. It explicitly checks for and prevents TPU runtime usage due to compatibility issues with modern JAX versions. Requires Python and pip. Outputs system information including NVIDIA-SMI details if a GPU is present.

```python
#@markdown # Environment setup
#@markdown **IMPORTANT NOTE**: Modern jax (>0.4) does not support the Colab TPU
#@markdown anymore, so don't select TPU runtime here. CPU and GPU work and are both fast enough.

# Install the right jax version for TPU/GPU/CPU
import os
if 'COLAB_TPU_ADDR' in os.environ:
  raise "TPU colab not supported."
elif 'NVIDIA_PRODUCT_NAME' in os.environ:
  !nvidia-smi
import jax
jax.devices()


# Get latest version of big_vision codebase.
!git clone --quiet --branch=main --depth=1 https://github.com/google-research/big_vision
!cd big_vision && git pull --rebase --quiet
!pip -q install -r big_vision/big_vision/requirements.txt
# Gives us ~2x faster gsutil cp to get the model checkpoints.
!pip3 -q install --no-cache-dir -U crcmod

%cd big_vision


import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

import jax
import jax.numpy as jnp
import ml_collections

from google.colab.output import _publish as publish

```

--------------------------------

### Example Model Parameter Overview Output

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/finetune_paligemma.ipynb

This snippet shows an example output of the `parameter_overview` function when executed. It lists various parameters within a model, detailing their hierarchical path, tensor shape, and data type (e.g., float16, float32). This output is crucial for understanding the model's structure and memory footprint.

```text
 == Model params == 
img/Transformer/encoder_norm/bias                                                (1152,)                float16
img/Transformer/encoder_norm/scale                                               (1152,)                float16
img/Transformer/encoderblock/LayerNorm_0/bias                                    (27, 1152)             float16
img/Transformer/encoderblock/LayerNorm_0/scale                                   (27, 1152)             float16
img/Transformer/encoderblock/LayerNorm_1/bias                                    (27, 1152)             float16
img/Transformer/encoderblock/LayerNorm_1/scale                                   (27, 1152)             float16
img/Transformer/encoderblock/MlpBlock_0/Dense_0/bias                             (27, 4304)             float16
img/Transformer/encoderblock/MlpBlock_0/Dense_0/kernel                           (27, 1152, 4304)       float16
img/Transformer/encoderblock/MlpBlock_0/Dense_1/bias                             (27, 1152)             float16
img/Transformer/encoderblock/MlpBlock_0/Dense_1/kernel                           (27, 4304, 1152)       float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/bias             (27, 16, 72)           float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/key/kernel           (27, 1152, 16, 72)     float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/bias             (27, 1152)             float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/out/kernel           (27, 16, 72, 1152)     float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/bias           (27, 16, 72)           float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/query/kernel         (27, 1152, 16, 72)     float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/bias           (27, 16, 72)           float16
img/Transformer/encoderblock/MultiHeadDotProductAttention_0/value/kernel         (27, 1152, 16, 72)     float16
img/embedding/bias                                                               (1152,)                float16
img/embedding/kernel                                                             (14, 14, 3, 1152)      float16
img/head/bias                                                                    (2304,)                float16
img/head/kernel                                                                  (1152, 2304)           float16
img/pos_embedding                                                                (1, 256, 1152)         float16
llm/embedder/input_embedding                                                     (257152, 2304)         float16
llm/final_norm/scale                                                             (2304,)                float16
llm/layers/attn/attn_vec_einsum/w                                                (26, 8, 256, 2304)     float32
llm/layers/attn/kv_einsum/w                                                      (26, 2, 4, 2304, 256)  float32
llm/layers/attn/q_einsum/w                                                       (26, 8, 2304, 256)     float32
llm/layers/mlp/gating_einsum                                                     (26, 2, 2304, 9216)    float16
llm/layers/mlp/linear                                                            (26, 9216, 2304)       float16
llm/layers/post_attention_norm/scale                                             (26, 2304)             float16
llm/layers/post_ffw_norm/scale                                                   (26, 2304)             float16
llm/layers/pre_attention_norm/scale                                             (26, 2304)             float16
llm/layers/pre_ffw_norm/scale                                                    (26, 2304)             float16
```

--------------------------------

### Render Training Examples In-line

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/finetune_paligemma.ipynb

Generates HTML to display training examples, combining images and their corresponding captions. It includes helper functions to convert images to base64 encoded JPEGs for in-line display and to format the HTML for each example. This is useful for visualizing model inputs and outputs during debugging.

```python
# @title Inspect training examples.
def render_inline(image, resize=(128, 128)):
  """Convert image into inline html."""
  image = Image.fromarray(image)
  image.resize(resize)
  with io.BytesIO() as buffer:
    image.save(buffer, format='jpeg')
    image_b64 = str(base64.b64encode(buffer.getvalue()), "utf-8")
    return f"data:image/jpeg;base64,{image_b64}"

def render_example(image, caption):
  image = ((image + 1)/2 * 255).astype(np.uint8)  # [-1,1] -> [0, 255]
  return f"""
    <div style="display: inline-flex; align-items: center; justify-content: center;">
        <img style="width:128px; height:128px;" src="{render_inline(image, resize=(64,64))}" />
        <p style="width:256px; margin:10px; font-size:small;">{html.escape(caption)}</p>
    </div>
    """

html_out = ""
for idx, example in zip(range(8), train_data_iterator()):
  caption = postprocess_tokens(example["text"])  # detokenize model input.
  caption = caption[len("caption en\n"):]        # strip prefix
  html_out += render_example(example["image"], caption)

print("Training examples")
display(HTML(html_out))
```

--------------------------------

### Install Dependencies and Fetch big_vision Code (Python)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/finetune_paligemma.ipynb

This snippet fetches the big_vision repository using Git and installs necessary Python dependencies like 'overrides', 'ml_collections', 'einops', and 'sentencepiece'. It also includes a check to prevent usage with remote TPUs in Colab.

```python
# @title Fetch big_vision code and install dependencies.
import os
import sys

# TPUs with
if "COLAB_TPU_ADDR" in os.environ:
  raise "It seems you are using Colab with remote TPUs which is not supported."

# Append big_vision code to python import path
if "big_vision_repo" not in sys.path:
  sys.path.append("big_vision_repo")

# Install missing dependencies. Assume jax~=0.4.25 with GPU available.
!pip3 install -q "overrides" "ml_collections" "einops~=0.7" "sentencepiece"

```

--------------------------------

### Configure and Load Model using Python

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

This Python code loads a model configuration from a specified file and sets up arguments for a SigLip model. It imports necessary modules and prepares the configuration object for model initialization.

```python
files.view('big_vision/big_vision/configs/proj/image_text/siglip_lit_coco.py')
from big_vision.configs.proj.image_text import siglip_lit_coco as lit_coco
arg = 'txt=bert_base,img=B/16,img_head,init=LiT-B16B.npz'
config = lit_coco.get_config(arg)
```

--------------------------------

### Configure TPU or GPU Acceleration

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

Sets up Colab TPUs if available, or checks for NVIDIA GPU availability using `nvidia-smi`. It then returns the list of available JAX devices, which will be either TPUs or GPUs.

```python
# Set up Colab  TPUs (if available).
import os
if 'COLAB_TPU_ADDR' in os.environ:
  import jax.tools.colab_tpu
  jax.tools.colab_tpu.setup_tpu()
else:
  !nvidia-smi
import jax
jax.devices()
```

--------------------------------

### Download Image and Display using Python

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

This snippet downloads an image if it doesn't exist and then loads it using PIL and NumPy for display with Matplotlib. It shows basic image handling and visualization.

```python
!test -f apple-ipod.jpg || wget https://cdn.openai.com/multimodal-neurons/assets/apple/apple-ipod.jpg

labels = [
    'an apple',
    'an ipod',
    'granny smith',
    'an apple with a note saying "ipod"',
    'an adversarial attack',
]

import PIL
import numpy as np
img = np.array(PIL.Image.open('apple-ipod.jpg'))
import matplotlib.pyplot as plt
plt.imshow(img)
img.shape, img.dtype
```

--------------------------------

### Example Big Vision Metrics Output Structure

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

This represents the expected output after successfully loading and parsing the metrics file. It's a Python dictionary containing various evaluation metrics, such as 'step', 'z/0shot/oxford_iiit_pet_accuracy', and 'z/secs/eval/disclf'. The values indicate performance or status at a given step.

```json
{
'step': 0,
'z/0shot/oxford_iiit_pet_accuracy': 0.8103025347506132,
'z/secs/eval/disclf': 73.91012993699997
}
```

--------------------------------

### Download Model Checkpoint using gsutil

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

This command-line snippet downloads a model checkpoint file (LiT-B16B.npz) from a Google Cloud Storage bucket if it doesn't already exist. This is a prerequisite for loading pre-trained model weights.

```bash
!test -f LiT-B16B.npz || gsutil cp gs://vit_models/lit/LiT-B16B.* .
```

--------------------------------

### Set up Python Path and Auto-reloading

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

Configures the Python environment by adding the big_vision directory to sys.path, allowing it to be imported. It also enables autoreload for enhanced development workflow in interactive environments like Colab.

```python
import sys
bv_path = './big_vision'
if bv_path not in sys.path:
  sys.path.insert(0, bv_path)

%load_ext autoreload
%autoreload 2
```

--------------------------------

### Image Preprocessing with NaFlex

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP2_demo.ipynb

Demonstrates the NaFlex aspect-preserving preprocessing method for images. This example involves splitting an image into patches of a specified size using NaFlex.

```python
# Example document image with height >> width.
img = PIL.Image.open('siglip2_screenshot.jpg')
print(img.size)
# img  # show image in original resolution
```

```python
# Split the image into 1024 naflex patches of size 16.
pp_naflex = get_pp_naflex(1024, 16)
```

--------------------------------

### Initialize Model Parameters using Python

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

This snippet initializes model parameters using JAX. It imports the model dynamically based on the configuration, creates initial parameter arrays with specified shapes and data types, and then initializes the model with a random key.

```python
# Initialize template params...
import importlib
import jax.numpy as jnp

model_mod = importlib.import_module(f'big_vision.models.{config.model_name}')

model = model_mod.Model(**config.model)

init_params = [
    jnp.zeros(shape, dtype)
    for shape, dtype in zip(config.init_shapes, config.init_types)
]

params0 = model.init(jax.random.PRNGKey(42), *init_params)['params'].unfreeze()
```

--------------------------------

### Initialize CLIPPO Model and Load Parameters

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/clippo/clippo_colab.ipynb

This Python code imports necessary modules from Big Vision, defines the image resolution, loads the model configuration, dynamically imports the model module based on the config, and initializes the model. It then loads the pre-trained model parameters from the downloaded checkpoint file using a utility function. This sets up the model for inference or further training.

```python
from big_vision.configs.proj.clippo import train_clippo
from big_vision import utils

# The models are trained for resolution 224
RES = 224

# Load model module
config = train_clippo.get_config()
model_module = importlib.import_module(f'big_vision.models.{config.model_name}')
model = model_module.Model(**config.model)

# Load model parameters
params = utils.load_checkpoint(None, checkpoint_path)['params']
```

--------------------------------

### Clone and Setup Repository (Python)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/givt/givt_demo_colab.ipynb

This snippet clones the Big Vision GitHub repository, pulls the latest changes, and installs necessary dependencies including specific versions of Keras, TensorFlow, and TensorFlow Probability. It also installs crcmod for accelerating gsutil operations and changes the working directory to 'big_vision'.

```python
#@markdown Clone and set up repository
!git clone --branch=main --depth=1 https://github.com/google-research/big_vision
!cd big_vision && git pull

# Install dependencies - pin TensorFlow-related packages to ensure compatibility
# which might not be needed in in the future
!echo -e "keras==3.0.5\ntensorflow==2.16.1\ntensorflow-probability==0.24.0" > big_vision/big_vision/constraints.txt
!pip install -r big_vision/big_vision/requirements.txt -c big_vision/big_vision/constraints.txt

# To fix/accelerate gsutil cp
!pip3 -q install --no-cache-dir -U crcmod

%cd big_vision
```

--------------------------------

### Set up Distributed Training on Cloud TPU (Bash)

Source: https://context7.com/google-research/big_vision/llms.txt

Provides bash commands for setting up a multi-host Cloud TPU VM for distributed training. It includes steps for creating a TPU VM, cloning the Big Vision repository, copying the code to all hosts, and executing a setup script.

```bash
# Create TPU VM with 32 cores (4 hosts)
export NAME=my-tpu-machine
export ZONE=europe-west4-a
export GS_BUCKET_NAME=my_bucket

gcloud compute tpus tpu-vm create $NAME \
    --zone $ZONE \
    --accelerator-type v3-32 \
    --version tpu-ubuntu2204-base

# Install big_vision on all hosts
git clone https://github.com/google-research/big_vision
gcloud compute tpus tpu-vm scp --recurse big_vision/big_vision $NAME: \
    --zone=$ZONE --worker=all
gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all \
    --command "bash big_vision/run_tpu.sh"

```

--------------------------------

### Clone and Install Big Vision Python Packages

Source: https://github.com/google-research/big_vision/blob/main/README.md

Clones the Big Vision repository and installs necessary Python dependencies using pip. It's recommended to use a virtual environment. This process involves updating pip, installing from a requirements file, and optionally installing the JAX library with CUDA support.

```bash
git clone https://github.com/google-research/big_vision
cd big_vision/
pip3 install --upgrade pip
pip3 install -r big_vision/requirements.txt
pip3 install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```

--------------------------------

### Install crcmod Package

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP2_demo.ipynb

Installs the 'crcmod' package using pip3 for accelerated gsutil copy operations. This package is crucial for improving the performance of model checkpoint downloads.

```shell
!pip3 -q install --no-cache-dir -U crcmod
```

--------------------------------

### Load UViM Model Checkpoints

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/uvim_color_task.ipynb

This code loads pre-trained model weights for both the stage-I (oracle) and stage-II (language model) components of the UViM model. It uses gsutil for downloading from cloud storage and JAX's device_put for efficient data placement on accelerators.

```python
# Load checkpoints
!gsutil cp -n gs://big_vision/uvim/color_stageI_params.npz gs://big_vision/uvim/color_stageII_params.npz .

oracle_params, oracle_state = vit.load(None, "color_stageI_params.npz")
oracle_params = jax.device_put({"params": oracle_params, "state": oracle_state})

lm_params = vtt.load(None, "color_stageII_params.npz")
lm_params = jax.device_put({"params": lm_params})
```

--------------------------------

### Configure and Load SigLip2 Model Checkpoint

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP2_demo.ipynb

This snippet configures the SigLip2 model variant, resolution, and embedding dimensions. It then downloads the corresponding checkpoint if it doesn't exist locally and initializes the model using the specified configuration. The model is loaded with parameters from the downloaded checkpoint.

```python
VARIANT, RES = 'g-opt/16', 384

CKPT = f'siglip2_{VARIANT.lower().replace("/", "")}_{RES}.npz'
TXTVARIANT, PATCH_SIZE = VARIANT.split('/')
EMBDIM = {'B': 768, 'L': 1024, 'So400m': 1152, 'g-opt': 1536}[TXTVARIANT]
# Note: The g-opt vision encoder is paired with a So400m text encoder
TXTVARIANT = 'So400m' if TXTVARIANT == 'g-opt' else TXTVARIANT
PATCH_SIZE = int(PATCH_SIZE)
VOCAB = 256_000
SEQLEN = 64

# It is significantly faster to first copy the checkpoint (30s vs 8m30 for B and 1m vs ??? for L)
!test -f /tmp/{CKPT} || gsutil cp gs://big_vision/siglip2/{CKPT} /tmp/

import big_vision.models.proj.image_text.two_towers as model_mod

model_cfg = ml_collections.ConfigDict(dict(
    image_model='vit',
    image=dict(
        pool_type='map',
        scan=True,
        variant=VARIANT,
    ),
    text_model='proj.image_text.text_transformer',
    text=dict(
        scan=True,
        variant=TXTVARIANT,
        vocab_size=256_000,
    ),
    out_dim=[None, EMBDIM],
    bias_init=-10,  # without this arg, no "b" param is added
))
if RES == 'naflex':
  model_cfg.image_model = 'proj.image_text.naflex_vit'
  model_cfg.image = dict(
      pool_type='map',
      nposemb=16,
      posemb='learn_2d',
      variant='B'
  )

model = model_mod.Model(**model_cfg)

# Using `init_params` is slower but will lead to `load` below performing sanity-checks.
# init_params = jax.jit(model.init, backend="cpu")(jax.random.PRNGKey(42), jnp.zeros([1, RES, RES, 3], jnp.float32), jnp.zeros([1, SEQLEN], jnp.int32))['params']
init_params = None  # Faster but bypasses loading sanity-checks.

params = model_mod.load(init_params, f'/tmp/{CKPT}', model_cfg)
```

--------------------------------

### Run Inference and Render Examples (Python/JAX/Matplotlib)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/uvim_panoptic_task.ipynb

This Python script demonstrates how to run the trained UVIM model on a few examples from the COCO validation dataset. It iterates through the dataset, performs predictions using the loaded model parameters, and visualizes the results (original image and predicted segmentation) using Matplotlib. It includes helper functions for rendering and processing.

```python
# Run the model in a few examples:
from matplotlib import pyplot as plt
from matplotlib import patches
from big_vision.trainers.proj.uvim import coco_utils

num_examples = 4
data = dataset.batch(1).take(num_examples).as_numpy_iterator()
key = jax.random.PRNGKey(0)
temperature = jnp.array(1e-7)

def render_example(image, prediction, with_legend=True):
  f, ax = plt.subplots(1, 2, figsize=(10, 10))
  ax[0].imshow(image*0.5 + 0.5)
  ax[0].axis("off")

  rgb, info = coco_utils.rgb_panoptic_from_twochannels(prediction, boundaries=True)
  ax[1].matshow(rgb)
  ax[1].axis("off")

  if with_legend:
    handles = []
    for instance in info.values():
      handles.append(patches.Patch(
          facecolor=np.array(instance["color"])/255.0,
          edgecolor='black', label=instance["name"]))
    ax[1].legend(handles=handles, loc=(1.04, 0.0));


for idx, batch in enumerate(data):
  subkey = jax.random.fold_in(key, idx)
  code = predict_code(lm_params, batch, key, temperature)
  aux_inputs = task.input_pp(batch, config.oracle)
  prediction = code2labels(oracle_params, code, aux_inputs["ctx"])
  render_example(batch["image"][0], prediction[0])
```

--------------------------------

### Invoke Stage I Panoptic Training (UViM VQ-VAE)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/README.md

Initiates the stage I panoptic training for the UViM model using a VQ-VAE architecture. It specifies the configuration file and sets a working directory based on the current date and time. The ':singlehost' parameter indicates a lightweight configuration for a single host.

```python
python3 -m big_vision.trainers.proj.uvim.vqvae --config big_vision/configs/proj/uvim/vqvae_coco_panoptic.py:singlehost --workdir workdirs/`date '+%m-%d_%H%M'`
```

--------------------------------

### Configure Kaggle Credentials (Python)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/finetune_paligemma.ipynb

Sets up environment variables for Kaggle username and API key using Colab's `userdata.get` API. It also configures `XLA_PYTHON_CLIENT_MEM_FRACTION` to preallocate memory, crucial for avoiding Out-Of-Memory errors during model finetuning on memory-constrained environments like the T4 runtime.

```python
import os
from google.colab import userdata

# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env
# vars as appropriate or make your credentials available in ~/.kaggle/kaggle.json

os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')

# The T4 runtime is tight on memory to finetune this model. Preallocate
# all memory ahead of time to avoid OOM'ing due to fragmentation.
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.0"

```

--------------------------------

### Clone and Update big_vision Repository

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

Clones the big_vision repository from GitHub and pulls the latest changes. This is a prerequisite for using the big_vision codebase. It ensures you have the most up-to-date version of the library.

```bash
!git clone --branch=main --depth=1 https://github.com/google-research/big_vision
!cd big_vision && git pull
```

--------------------------------

### Inspect Evaluators

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

Retrieves and displays the list of pre-defined evaluators available within the configuration. The `set_max_height(222)` function might be used to control the display or processing height, and `config.evals` accesses the evaluators.

```python
# From all the pre-defined evaluators...
set_max_height(222)
config.evals
```

--------------------------------

### Create Evaluation/Inference Prediction Loop with JAX

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/finetune_paligemma.ipynb

Generates model predictions for a given data iterator. It handles batching, padding, and tokenization. The function can limit the number of examples processed and supports different sampling strategies.

```python
import jax
import jax.numpy as jnp
import numpy as np

# Assuming 'SEQLEN', 'params', 'data_sharding', 'big_vision', 'decode', 'postprocess_tokens' are defined elsewhere.

def make_predictions(data_iterator, *, num_examples=None,
                     batch_size=4, seqlen=SEQLEN, sampler="greedy"):
  outputs = []
  while True:
    examples = []
    try:
      for _ in range(batch_size):
        examples.append(next(data_iterator))
        examples[-1]["_mask"] = np.array(True)
    except StopIteration:
      if len(examples) == 0:
        return outputs

    while len(examples) % batch_size:
      examples.append(dict(examples[-1]))
      examples[-1]["_mask"] = np.array(False)

    batch = jax.tree.map(lambda *x: np.stack(x), *examples)
    batch = big_vision.utils.reshard(batch, data_sharding)

    tokens = decode({"params": params}, batch=batch,
                    max_decode_len=seqlen, sampler=sampler)

    tokens, mask = jax.device_get((tokens, batch["_mask"]))
    tokens = tokens[mask]
    responses = [postprocess_tokens(t) for t in tokens]

    for example, response in zip(examples, responses):
      outputs.append((example["image"], response))
      if num_examples and len(outputs) >= num_examples:
        return outputs
```

--------------------------------

### Run UViM Model Prediction and Rendering

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/uvim_color_task.ipynb

This code demonstrates running the UViM model on sample images from the prepared dataset. It iterates through a batch of images, generates color predictions using the loaded model and parameters, and then visualizes the original and colorized images using Matplotlib.

```python
# Run the model in a few examples:
from matplotlib import pyplot as plt

num_examples = 4
data = dataset.batch(1).take(num_examples).as_numpy_iterator()
key = jax.random.PRNGKey(0)
temperature = jnp.array(1.0)

def render_example(image, prediction):
  f, ax = plt.subplots(1, 2, figsize=(10, 10))
  ax[0].imshow(image*0.5 + 0.5)
  ax[0].axis("off")
  ax[1].imshow(prediction*0.5 + 0.5)
  ax[1].axis("off")

for idx, batch in enumerate(data):
  subkey = jax.random.fold_in(key, idx)
  code = predict_code(lm_params, batch, key, temperature)
  aux_inputs = task.input_pp(batch, config.oracle)
  prediction = code2labels(oracle_params, code, aux_inputs["ctx"])
  render_example(batch["image"<ctrl63>
```

--------------------------------

### Load Model Checkpoints (Python/JAX)

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/uvim/uvim_panoptic_task.ipynb

This Python snippet loads pre-trained model checkpoints for the UVIM panoptic stage-I and stage-II models from Google Cloud Storage. It utilizes JAX for loading and device placement, preparing the models for inference.

```python
# Load checkpoints
!gsutil cp -n gs://big_vision/uvim/panoptic_stageI_params.npz gs://big_vision/uvim/panoptic_stageII_params.npz .

oracle_params, oracle_state = vit.load(None, "panoptic_stageI_params.npz")
oracle_params = jax.device_put({"params": oracle_params, "state": oracle_state})

lm_params = vtt.load(None, "panoptic_stageII_params.npz")
lm_params = jax.device_put({"params": lm_params})
```

--------------------------------

### Download and Unzip Unifont Files

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/clippo/clippo_colab.ipynb

This code snippet downloads two compressed Unifont files (unifont-9.0.06.hex.gz and unifont_upper-9.0.06.hex.gz) using wget and then unzips them using gunzip. Finally, it moves the unzipped files to a specified directory within the project structure. This is necessary for text rendering capabilities.

```shell
!wget https://unifoundry.com/pub/unifont/unifont-9.0.06/font-builds/unifont-9.0.06.hex.gz https://unifoundry.com/pub/unifont/unifont-9.0.06/font-builds/unifont_upper-9.0.06.hex.gz
!gunzip unifont-9.0.06.hex.gz unifont_upper-9.0.06.hex.gz
!mv unifont-9.0.06.hex unifont_upper-9.0.06.hex big_vision/big_vision/pp/proj/clippo/
```

--------------------------------

### Run Distillation Training Command

Source: https://context7.com/google-research/big_vision/llms.txt

Command-line instruction to execute the Knowledge Distillation training using the defined configuration file. It specifies the Python module to run, the configuration path, and the working directory for the experiment.

```bash
# python -m big_vision.trainers.proj.distill.distill \
#     --config big_vision/configs/proj/distill/bit_i1k.py \
#     --workdir gs://bucket/distill_resnet50
```

--------------------------------

### Define Colab Output Height Adjustment Function

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

Defines a Python function `set_max_height` that uses Javascript to adjust the maximum scrollable height of Colab output cells. This is useful for controlling the display of large outputs.

```python
from absl import flags
from absl import logging
import tensorflow_datasets as tfds
from google.colab import files

logging.set_verbosity(logging.INFO)

def set_max_height(max_height):
  """Limits scrollable area of output cell to `max_height` pixels."""
  import IPython.display
  IPython.display.display(IPython.display.Javascript('''
    google.colab.output.setIframeHeight(0, true, {maxHeight: %d})
  ''' % max_height))

```

--------------------------------

### JAX JIT-Compiled Prediction Function for Classification

Source: https://context7.com/google-research/big_vision/llms.txt

Defines a JAX JIT-compiled function to perform model predictions on batches of data. This is typically used during evaluation to get logits for calculating metrics.

```python
# Define prediction and loss functions
@jax.jit
def predict_fn(train_state, batch):
    logits, _ = model.apply(
        {'params': train_state['params']},
        batch['image'],
        train=False
    )
    return logits
```

--------------------------------

### Prepare Oxford-IIIT Pet Dataset

Source: https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb

Downloads and prepares the Oxford-IIIT Pet dataset using TensorFlow Datasets. This step ensures the dataset is available locally in the correct format for evaluation. It uses `tfds.builder` to access and manage the dataset.

```python
# Prepare pets dataset.
tfds.builder('oxford_iiit_pet').download_and_prepare()
```