Add files using upload-large-folder tool

Tipsomaly/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Alireza Salehi et al.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights  
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell      
+copies of the Software, and to permit persons to whom the Software is         
+furnished to do so, subject to the following conditions:                       
+
+The above copyright notice and this permission notice shall be included in    
+all copies or substantial portions of the Software.                           
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR    
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,      
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE   
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER        
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN     
+THE SOFTWARE.

Tipsomaly/README.md

@@ -0,0 +1,131 @@
+# Tipsomaly (ICASSP 2026)
+
+Official PyTorch implementation of [TIPS Over Tricks: Simple Prompts for Effective Zero-shot Anomaly Detection](https://arxiv.org/abs/2602.03594) — a spatially-aware zero-shot anomaly detection pipeline built on the [TIPS](https://arxiv.org/abs/2410.16512) vision-language model, using decoupled prompts and local evidence injection to improve image-level and pixel-level performance.
+
+---
+
+## Table of Contents
+
+- [📖 Introduction](#-introduction)
+- [📊 Results](#-results)
+- [🚀 Quickstart](#-quickstart)
+- [📍 Checkpoints](#-checkpoints)
+- [🔧 Setup](#-setup)
+- [🗂️ Datasets](#️-datasets)
+- [🛠️ Training](#️-training)
+- [⚖️ License](#️-license)
+- [🔗 Citation](#-citation)
+<!-- - [🙏 Acknowledgements](#-acknowledgements) -->
+
+---
+
+## 📖 Introduction
+Anomaly detection identifies departures from expected behavior in safety-critical settings. When target-domain normal data are unavailable, zero-shot anomaly detection (ZSAD) leverages vision-language models (VLMs). However, CLIP's coarse image-text alignment limits both localization and detection due to (i) spatial misalignment and (ii) weak sensitivity to fine-grained anomalies; prior work compensates with complex auxiliary modules yet largely overlooks the choice of backbone. We revisit the backbone and use TIPS-a VLM trained with spatially aware objectives. While TIPS alleviates CLIP's issues, it exposes a distributional gap between global and local features. We address this with decoupled prompts-fixed for image-level detection and learnable for pixel-level localization-and by injecting local evidence into the global score. Without CLIP-specific tricks, our TIPS-based pipeline improves image-level performance by 1.1-3.9% and pixel-level by 1.5-6.9% across seven industrial datasets, delivering strong generalization with a lean architecture.
+
+![Architecture](imgs/Models_Architecture_page-0001.jpg)
+
+---
+
+## 📊 Results
+Across 14 industrial and medical benchmarks, Tipsomaly consistently outperforms prior CLIP-based zero-shot methods while remaining lightweight. The figure below summarizes its performance across datasets.
+
+![Results](imgs/results-table.png)
+
+We compare pixel-level anomaly maps from Tipsomaly with prior CLIP-based methods (AdaCLIP and AnomalyCLIP) on industrial and medical samples. As shown in the following figure, Tipsomaly more accurately localizes anomalous regions across both domains.
+
+![Qualitative-results](imgs/Qualitative_results_page-0001.jpg)
+
+---
+
+## 🚀 Quickstart
+**Quick start (recommended):**  
+Use our ready-to-run [Kaggle notebook](https://www.kaggle.com/code/sepehrnoey/tipsomaly-reproduction) to reproduce the results with the provided checkpoints.
+
+## 📍 Checkpoints
+You can find our checkpoints trained on MVTec and VisA in the directory [workspaces](./workspaces/) or on google drive at [this link](https://drive.google.com/file/d/1yvgZYHFljFkGwTD0DYOdrZ_95ZbmjDcC/view?usp=sharing).
+
+---
+
+
+
+
+## 🔧 Setup
+To setup the environment and run experiments, follow these steps.
+1. Install the required dependencies with the following command.
+```
+pip install -r requirements.txt
+```
+2. Download your desired checkpoints of the TIPS model. For this purpose, edit and run the mentioned script to download model components. After downloading the components, place them in a desired directory.
+```
+bash model/tips/checkpoints/download_checkpoints.sh
+```
+```
+mkdir tips && mv tips_oss_l14_highres_distilled_text.npz tips_oss_l14_highres_distilled_vision.npz tokenizer.model tips/
+```
+3. Prepare your datasets in the format described in the [Datasets](#️-datasets) section and place them in a desired directory.
+```
+mkdir data-root && mv mvtec visa data-root/
+``` 
+4. Now, by setting the correct paths for the model, checkpoints, and the data root path, you can edit the runnable script `reproduce.sh` to run experiments. Here is a sample:
+```
+models_dir="path/to/tips"
+data_root_dir="path/to/data-root"
+model_version='l14h' 
+checkpoint_path="./workspaces/trained_on_mvtec_default/vegan-arkansas/checkpoints"
+
+python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path --data_root_dir $data_root_dir --dataset visa --epoch 2 --model_version $model_version --fixed_prompt_type industrial
+
+```
+
+## 🗂️ Datasets
+The datasets should be set in the following format. Each dataset, must contain a file named `meta.json` which contains information about the dataset. Common datasets for anomaly detection and instructions to generate `meta.json` files can be found in the [AdaCLIP GitHub repository](https://github.com/caoyunkang/AdaCLIP). The general expected structure of the datasets is like the following:
+```
+data-root/
+├── mvtec/
+│   ├── meta.json
+│   ├── bottle/
+│   │   ├── train/
+│   │   │   └── good/
+│   │   │       └── 000.png
+│   │   ├── test/
+│   ���   │   ├── good/
+│   │   │   │   └── 000.png
+│   │   │   ├── broken_large/
+│   │   │   │   └── 000.png
+│   │   │   └── broken_small/
+│   │   │       └── 000.png
+│   │   └── ground_truth/
+│   │       ├── broken_large/
+│   │       │   └── 000_mask.png
+│   │       └── broken_small/
+│   │           └── 000_mask.png
+│
+├── visa/
+├── mpdd/
+├── ....
+....
+```
+
+## 🛠️ Training  
+You can use our ready-to-run [Kaggle notebook](https://www.kaggle.com/code/sepehrnoey/tipsomaly-train) to train your model and test on your desired dataset. Steps to prepare the datasets and downloading the base model components can be found at the [Setup](#-setup) section. You can also find the commands to train the model and test multiple datasets in a loop in the [train_test.sh](train_test.sh)
+
+## 🔗 Citation
+If you find this project helpful for your research, please consider citing the following BibTeX entry.
+
+
+<!-- 📚 [Paper Link](https://arxiv.org/pdf/2504.11055) -->
+
+**BibTeX:**
+```bibtex
+@article{salehi2026tips,
+  title={TIPS Over Tricks: Simple Prompts for Effective Zero-shot Anomaly Detection},
+  author={Salehi, Alireza and Karami, Ehsan and Noey, Sepehr and Noey, Sahand and Yamada, Makoto and Hosseini, Reshad and Sabokrou, Mohammad},
+  journal={arXiv preprint arXiv:2602.03594},
+  year={2026}
+}
+```
+
+<!-- ## 🙏 Acknowledgements -->
+
+## ⚖️ License
+This project is licensed under the MIT License. Please refer to the [LICENSE](LICENSE) file for more information.

Tipsomaly/datasets/dataset.py

@@ -0,0 +1,115 @@
+import torch.utils.data as data
+import json
+import random
+from PIL import Image
+import numpy as np
+import os
+
+def sample_data(meta_info, k_shot):
+    sampled_data = []
+    complement_data = [] 
+    
+    for cls_name, data_list in meta_info.items():
+        nrm_smpls = [item for item in data_list if item['anomaly'] == 0]
+        anm_smpls = [item for item in data_list if item['anomaly'] == 1]
+        
+        if k_shot > 0:
+            n_samples = k_shot
+            n_nrm_smpls = min(int(n_samples / 2), len(nrm_smpls))
+            n_anm_smpls = min(int(n_samples / 2), len(anm_smpls))
+
+        cls_data = []
+        cls_data.extend(random.sample(nrm_smpls, n_nrm_smpls))
+        cls_data.extend(random.sample(anm_smpls, n_anm_smpls))
+        sampled_data.extend(cls_data)
+        
+        complement_class_data = [item for item in data_list if item not in cls_data]
+        complement_data.extend(complement_class_data)
+        
+        print(f'num samples for cls {cls_name}, norm: {n_nrm_smpls}, anom: {n_anm_smpls}')
+        
+    return sampled_data, complement_data
+
+class Dataset(data.Dataset):
+    def __init__(self, roots, transform, target_transform, kwargs=None):
+        self.roots = roots
+        self.transform = transform
+        self.target_transform = target_transform
+        split='test'
+        
+        meta_infos = {}
+        for root in roots:
+            with open(f'{root}/meta.json', 'r') as f:
+                meta_info = json.load(f)
+                for cls in meta_info[split]:
+                    meta_info[split][cls] = [{**s, 'root': root} for s in meta_info[split][cls]]
+                    
+                    if cls in meta_infos:
+                        meta_infos[cls].extend(meta_info[split][cls])
+                        meta_infos[cls].extend(meta_info[split][cls])
+                    else:
+                        meta_infos[cls] = meta_info[split][cls]
+                        
+        meta_info_classes = list(meta_infos.keys())
+
+        self.selected_class = kwargs.class_name
+        self.cls_names = meta_info_classes if self.selected_class == ['all'] else self.selected_class
+        
+        self.data_all = []
+        for cls_name in self.cls_names:
+            self.data_all.extend(meta_infos[cls_name])
+
+        self.dataset_name = kwargs.dataset
+        self.class_ids = list(range(len(self.cls_names)))
+        self.class_name_map_class_id = {k: index for k, index in zip(self.cls_names, self.class_ids)}
+
+        # Few-shot dataset (splitting...)
+        self.k_shot = kwargs.k_shot
+        if not self.k_shot == 0:
+            sampled_sets = sample_data(meta_infos, self.k_shot)
+            
+            if kwargs.type == 'train':
+                self.data_all = sampled_sets[0] 
+            
+            elif kwargs.type == 'test' and kwargs.train_dataset == self.dataset_name:
+                self.data_all = sampled_sets[1] 
+        
+        self.length = len(self.data_all)
+        print(f"number of dataset samples: {self.length}")
+            
+    def _process_image(self, data):
+        img_path, mask_path, cls_name, specie_name, anomaly = data['img_path'], data['mask_path'], \
+                                                                data['cls_name'], data['specie_name'], data['anomaly']
+        
+        root = data['root']    
+        full_img_path = os.path.join(root, img_path)
+        img = Image.open(full_img_path).convert('RGB')  
+
+        if anomaly == 0 or (not os.path.isfile(os.path.join(root, mask_path))):
+            img_mask = Image.fromarray(np.zeros((img.size[1], img.size[0])), mode='L')
+        else:
+            img_mask = np.array(Image.open(os.path.join(root, mask_path)).convert('L')) > 0
+            img_mask = Image.fromarray(img_mask.astype(np.uint8) * 255, mode='L')
+
+        img = self.transform(img) if self.transform is not None else img
+        img_mask = self.target_transform(img_mask)
+        img_mask = np.where(img_mask > 0.5, 1, 0)
+
+        result = {
+            'img': img,
+            'abnorm_mask': img_mask,
+            'cls_name': cls_name,
+            'anomaly': anomaly,
+            'img_path': os.path.join(root, img_path),
+            "cls_id": self.class_name_map_class_id[cls_name]
+        }
+
+        return result
+
+    def __len__(self):
+        return self.length
+    
+    def __getitem__(self, index):
+        data = self.data_all[index]
+        result = self._process_image(data)
+        return result

Tipsomaly/datasets/input_transforms.py

@@ -0,0 +1,23 @@
+from torchvision import transforms
+
+IMAGE_MEAN = (0, 0, 0)
+IMAGE_STD = (1.0, 1.0, 1.0)
+
+class Ensure3Channels:
+    def __call__(self, img):
+        return img.convert('RGB')  
+
+def create_transforms_tips(image_size):
+    transform = transforms.Compose([
+        Ensure3Channels(),
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(IMAGE_MEAN, IMAGE_STD),
+    ])
+
+    target_transform = transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+    ])
+
+    return transform, target_transform

Tipsomaly/model/big_vision/configs/proj/paligemma/transfers/vatex_cap.py

@@ -0,0 +1,210 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to VATEX Video captioning.
+
+IMPORTANT: This config is based on an unreleased version of DeepMind Video
+Readers (DMVR). Users can either set up DMVR using the open source code from
+GitHub (see below for details), or add their own data loader of choice.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+TEXT_LEN = 64
+DATASET_NAME = 'vatex'
+# Numbers might need to be updated due to wipeout. Current from 2024-04-28
+SPLIT_SIZE = {'train': 22315, 'valid': 2584, 'test': 5135}
+
+
+def training_data(res, *, final_split, num_frames=8, stride=None):
+  """Creates training data config.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: Train on all train+valid data.
+    num_frames: number of sampled frames per video.
+    stride: stride at which the frames are sampled.
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  pp = '|'.join([
+      # prepare the frames by decoding, resizing, replicating, sampling:
+      f'video_decode({res})|video_replicate_img({num_frames},{num_frames})',
+      f'video_ensure_shape("image", {(num_frames, res, res, 3)})',
+      # pick one caption at random during training
+      'strfmt("caption en", outkey="prefix")',
+      'video_choice(inkey="caption/string", outkey="suffix")',
+      combine_and_keep_train(TEXT_LEN),
+  ])
+
+  c = bvcc.parse_arg('')
+  c.data = {}
+  splits = ['train', 'valid'] if final_split else ['train']
+  raise NotImplementedError('Please implement a video reader of choice!')
+  # For example DMVR https://github.com/google-deepmind/dmvr
+  # The reader should support the following arguments:
+  # - name: Name of the reader.
+  # - dataset_name: Name of the data set.
+  # - split: Data set split.
+  # - num_frames: Number of frames sampled from the video.
+  # - stride: Stride at which the video frames are sampled.
+  # - deterministic_fs: Whether to sample the frames starting at the first
+  #   frame or whether an offest should be chosen at random (if there are more
+  #   frames than num_frames * stride)
+  # - first_k_shards: Whether to only use the first k shards of the data
+  #   (optional but useful for speeding up intermediate evaluations).
+  for split in splits:
+    c.data[split] = SPLIT_SIZE[split]
+    c[split] = {'pp': pp}
+    c[split].data = dict(
+        # PLEASE ADD YOUR READER HERE:
+        name='<add_your_data_loader_here>',
+        dataset_name=DATASET_NAME, split=split,
+        num_frames=num_frames, stride=stride,
+        deterministic_fs=False)
+  return c
+
+
+def add_eval(c, res, num_frames=8, stride=None):
+  """Captioning evaluator."""
+  c_train = training_data(res, final_split=True, num_frames=num_frames, stride=stride)
+
+  pp = '|'.join([
+      # prepare the frames by decoding, resizing, replicating, sampling:
+      f'video_decode({res})|video_replicate_img({num_frames},{num_frames})',
+      f'video_ensure_shape("image", {(num_frames, res, res, 3)})',
+      'strfmt("caption en", outkey="prefix")',
+      'copy("example/video_id", "image/id")',
+      'copy("caption/string", "captions")',
+      combine_and_keep_eval(TEXT_LEN, keep=('image/id', 'captions')),
+  ])
+
+  for freq, name, split, first_k_shards, skip_first_eval in [
+      (1/8, 'minitrain', 'train', 2, False),  # To gauge memorization.
+      (1/4, 'minival', 'valid', 2, False),  # To monitor val progress.
+      (1, 'val', 'valid', None, False),  # To tune hparams.
+      (1, 'eval', 'test', None, False),  # final metric
+  ]:
+    c.evals[f'{DATASET_NAME}/{name}'] = dict(
+        type='proj.paligemma.transfers.coco_caption',
+        pred='decode', pred_kw={'max_decode_len': TEXT_LEN},
+        data={**c_train.train.data, 'split': split,
+              'first_k_shards': first_k_shards,
+              'deterministic_fs': True},
+        log_percent=freq, tokenizer=TOKENIZER,
+        pp_fn=pp, skip_first=skip_first_eval)
+
+
+def add_eval_pplx(c, res, num_frames=8, stride=None):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, final_split=True, num_frames=num_frames, stride=stride)
+
+  for name, split, first_k_shards in [
+      ('minitrain', 'train', 2),  # To gauge memorization.
+  ]:
+    c.evals[f'{DATASET_NAME}/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=1/8,  # Not too cheap, do 10x per run.
+        data={**c_train.train.data, 'split': split,
+              'first_k_shards': first_k_shards,
+              'deterministic_fs': True},
+        pp_fn=c_train.train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  add(lr=3e-6, wd=3e-7, total_epochs=10, **bvcc.arg(res=224, **c))
+
+
+sweep = sweep_best
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', num_frames=16, stride=7, res=224,
+                     freeze_vit=False, freeze_llm=False, final_split=False)
+
+  c.input = training_data(
+      c.res, final_split=c.final_split,
+      num_frames=c.num_frames, stride=c.stride)
+
+  c.total_epochs = 3
+  c.input.batch_size = 128
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 3e-6
+  c.wd = 3e-7
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+
+  # Learning-rate schedule.
+  sched = dict(decay_type='cosine', warmup_percent=0.05)
+  c.schedule = [
+      ('img/.*', None if c.freeze_vit else sched),
+      ('llm/.*', None if c.freeze_llm else sched),
+  ]
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, c.num_frames, c.stride)
+  add_eval_pplx(c, c.res, c.num_frames, c.stride)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  for split in c.input.data.keys():
+    c.input[split].shuffle_buffer_size = 10_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops',
+                  'proj.paligemma.video']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    for split in c.input.data.keys():
+      c.input[split].shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.first_k_shards = 1
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('minitrain', 'minival', 'val', 'eval'):
+    m.append((f'{DATASET_NAME}/{split}/cider'))
+  for split in ('minitrain', 'minival'):
+    m.append((f'{DATASET_NAME}/{split}/pplx/avg'))
+  return m
+

Tipsomaly/model/big_vision/configs/proj/paligemma/transfers/vizwizvqa.py

@@ -0,0 +1,160 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to VQAv2.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+
+def training_data(res, final_split, text_len=48):
+  """Creates training data config.
+
+  See (internal link)
+  You can add more arguments beside `res`, but give them good defaults.
+
+  Args:
+    res: The requested image resolution (eg 224)
+    final_split: Train on combined train+val
+    text_len: sequence length
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='vizwizvqa',
+      split='train+val' if final_split else 'train',
+  )
+  c.pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'choice_no_replacement(inkey="answers", outkey="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  return c
+
+
+def add_eval(c, res, text_len=48, **kw):
+  """VQAv2 evaluators."""
+  pp = '|'.join([
+      f'decode|resize({res})|value_range(-1, 1)',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'copy("image/filename", "question_id")',
+      combine_and_keep_eval(text_len, keep=('answers', 'question_id')),
+  ])
+
+  for freq, name, split in [
+      (1/8, 'minitrain', 'train[:5120]'),  # To gauge memorization. 400s on 32v2
+      (0.1, 'minival', 'val'),  # To tune hparams. 4k samples, full eval is fine.
+      (1.0, 'test', 'test'),  # For the test-server. SLOW.
+  ]:
+    c.evals[f'vizwizvqa/{name}'] = dict(
+        type='proj.paligemma.transfers.vqa',
+        pred='decode', pred_kw={'max_decode_len': text_len},
+        outfile=f'{{workdir}}/vizwiz_{name}.json',
+        out_question_key='image',
+        data={**training_data(res, True, text_len).data, 'split': split},
+        log_percent=freq, tokenizer=TOKENIZER, pp_fn=pp)
+    c.evals[f'vizwizvqa/{name}'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=48):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, True, text_len)  # Use mostly same settings as training.
+  for name, split in [
+      ('minitrain', 'train'),  # To gauge memorization.
+      ('minival', 'val'),  # To tune hparams
+  ]:
+    c.evals[f'vizwizvqa/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=1/8,
+        data={**c_train.data, 'split': split},
+        pp_fn=c_train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  add(**bvcc.arg(res=224, **c))
+  add(**bvcc.arg(res=448, **c))
+
+
+sweep = sweep_best  # Choose which sweep to run.
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', res=224, final_split=False)
+
+  c.input = training_data(c.res, c.final_split)
+
+  # Instead of epochs, you can also use `total_examples` or `total_steps`.
+  c.total_epochs = 10
+  c.input.batch_size = 256
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 0.00001
+  c.wd = 0.0
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+  c.schedule = dict(decay_type='cosine', warmup_percent=0.05)
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, batch_size=1024)
+  add_eval_pplx(c, c.res)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  # These probably do not need any change/tuning
+  c.input.shuffle_buffer_size = 25_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    c.input.shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.split = ev.data.split.split('[')[0] + '[:16]'
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('minival', 'minitrain'):
+    m.append(f'vizwizvqa/{split}/acc')
+    m.append(f'vizwizvqa/{split}/pplx/avg')
+  return m

Tipsomaly/model/big_vision/configs/proj/paligemma/transfers/widgetcap.py

@@ -0,0 +1,180 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to widgetcap (bbox drawn in the picture).
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+
+def training_data(res, *, final_split, text_len=32):
+  """Creates training data config.
+
+  See (internal link)
+  You can add more arguments beside `res`, but give them good defaults.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: Train on all train+dev data.
+    text_len: The max text length.
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='widgetcap',
+      split='train+dev' if final_split else 'train',
+  )
+  c.pp = '|'.join([
+      'decode',
+      f'resize({res}, antialias=True)',
+      'draw_bbox',
+      'value_range(-1, 1)',
+      'strfmt("caption en", outkey="prefix")',
+      'choice_no_replacement(inkey="texts", outkey="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  return c
+
+
+def add_eval(c, res, text_len=32, **kw):
+  """Captioning evaluator with cider/bleu-4/meteor/rouge/spice metrics."""
+  # Input eval pp without ground truth text and random crop.
+  pp_eval = '|'.join([
+      'copy("texts", "captions")',  # GT for evaluator.
+      'decode',
+      f'resize({res}, antialias=True)',
+      'draw_bbox',
+      'value_range(-1, 1)',
+      'strfmt("caption en", outkey="prefix")',
+      combine_and_keep_eval(text_len, keep=('image/id', 'captions')),
+  ])
+
+  for name, split in [
+      ('val', 'dev'),
+      ('eval', 'test'),
+  ]:
+    c.evals[f'widgetcap/{name}'] = dict(
+        type='proj.paligemma.transfers.coco_caption',
+        pred='decode', pred_kw={'max_decode_len': text_len},
+        data=dict(
+            name='widgetcap',
+            split=split,
+        ),
+        log_percent=0.1, tokenizer=TOKENIZER, pp_fn=pp_eval)
+    c.evals[f'widgetcap/{name}'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=32):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, final_split=True, text_len=text_len)  # Use mostly same settings as training.
+  for name, split in [
+      ('minitrain', 'train[:5%]'),  # To gauge memorization.
+      ('minival', 'dev'),  # To tune hparams.
+      ('eval', 'test'),  # To compute final publishable scores.
+  ]:
+    c.evals[f'widgetcap/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=0.05,  # Eval ~20x per run; it's cheap.
+        data={**c_train.data, 'split': split},
+        pp_fn=c_train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  # Based on sweeps (internal link) (widgetcap/val/cider).
+  # NOTE: dropout always on, see get_config.
+  add(lr=3e-6, wd=3e-7, total_epochs=4, **bvcc.arg(res=224, **c))
+  add(lr=3e-6, wd=3e-7, total_epochs=4, **bvcc.arg(res=448, **c))
+  # Not better: add(lr=3e-6, wd=3e-7, total_epochs=4, **bvcc.arg(res=896, **c))
+
+
+sweep = sweep_best  # Choose which sweep to run.
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', res=224, final_split=False)
+
+  c.input = training_data(c.res, final_split=c.final_split)
+
+  # Instead of epochs, you can also use `total_examples` or `total_steps`.
+  c.total_epochs = 4
+  c.input.batch_size = 64
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 3e-6
+  c.wd = 3e-7
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.1
+  c.schedule = dict(decay_type='cosine', warmup_percent=0.05)
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, batch_size=1024)
+  add_eval_pplx(c, c.res)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.1)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  # These probably do not need any change/tuning
+  c.input.shuffle_buffer_size = 50_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = [
+      'ops_general',
+      'ops_image',
+      'ops_text',
+      'proj.paligemma.ops',
+      'proj.paligemma.widgetcap',
+  ]
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    c.input.shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.split = ev.data.split.split('[')[0] + '[:16]'
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  # This function defines the default flatboard. If you want, it can be a lot
+  # fancier too, but the simplest way is a list of metric names.
+  m = ['training_loss']
+  for split in ('eval', 'minival', 'minitrain'):
+    m.append(f'widgetcap/{split}/pplx/avg')
+  for split in ('val', 'eval'):
+    m.append(f'widgetcap/{split}/cider')
+  return m

Tipsomaly/model/big_vision/input_pipeline.py

@@ -0,0 +1,357 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""ImageNet input pipeline."""
+import collections
+import functools
+import itertools
+import math
+import multiprocessing.pool
+
+from absl import logging
+from big_vision.datasets import sequence_packing
+import big_vision.datasets.core as ds_core
+import big_vision.pp.builder as pp_builder
+import big_vision.utils as u
+import einops
+import jax
+import numpy as np
+import tensorflow as tf
+
+
+DEFAULT_NUM_PARALLEL_CALLS = 100
+
+
+def make_for_train(
+    data, preprocess_fn, batch_size,
+    shuffle_buffer_size=None, cache_raw=False,
+    num_parallel_calls=DEFAULT_NUM_PARALLEL_CALLS, prefetch=2,
+    *,
+    pre_filter_fn=None, post_filter_fn=None,
+    pack=None, skip_errors=False,
+):
+  """Makes an input pipeline for training."""
+  # Use data filtering at your own risk: the actual split sizes won't be known
+  # in advance, so epoch-based things won't work correctly.
+
+  data = _add_tpu_host_options(data)
+
+  data = data.filter(pre_filter_fn) if pre_filter_fn else data
+  data = data.cache() if cache_raw else data
+
+  # First shuffle and then repeat (each with a different shuffle). This way
+  # the data for one epoch is all seen before the next one is processed and
+  # significantly affects the number of times each example is seen when
+  # processing for small number of epochs.
+  if shuffle_buffer_size:
+    data = data.shuffle(shuffle_buffer_size, reshuffle_each_iteration=True)
+  data = data.repeat(None)
+
+  data = data.map(preprocess_fn, num_parallel_calls=num_parallel_calls)
+  data = data.filter(post_filter_fn) if post_filter_fn else data
+
+  data = data.ignore_errors(log_warning=True) if skip_errors else data
+
+  if pack:
+    data = sequence_packing.pack_dataset(
+        data,
+        batch_size // jax.process_count() if batch_size else None,
+        pack.to_dict())
+
+  # Drop remainder makes shape fully static, so we can later use it if needed.
+  if batch_size:
+    data = data.batch(batch_size // jax.process_count(), drop_remainder=True)
+  if prefetch:  # None means autotune, but we never want that.
+    data = data.prefetch(prefetch)
+  return data
+
+
+def training(input_config):
+  """Reads the data from a single dataset, or mixes it from multiple.
+
+  The data is read either from one or mixed from multiple datasets, depending
+  on the `input_config`.
+
+  Args:
+    input_config: Configures the input pipeline. See input_pipeline_test for
+      examples.
+
+  Returns:
+    A tuple containing (possibly mixed) tf.data.Dataset and a total number of
+    training examples.
+  """
+  per_pipeline_configs = (
+      "shuffle_buffer_size", "cache_raw", "num_parallel_calls",
+      "pre_filter_fn", "post_filter_fn", "pack", "skip_errors")
+  def config_to_kw(config):
+    assert "filter_fn" not in config, "Deprecated; use `pre_filter_fn` instead."
+    return {k: config[k] for k in per_pipeline_configs if k in config}
+
+  batch_size = input_config.batch_size
+  # Handle separately the common case when no mixing happens.
+  if isinstance(input_config.data.get("name"), str):
+    train_data = ds_core.get(**input_config.data)
+    train_ds = make_for_train(
+        data=train_data.get_tfdata(ordered=False,
+                                   **input_config.get("tfdata", {})),
+        batch_size=batch_size,
+        preprocess_fn=pp_builder.get_preprocess_fn(input_config.get("pp")),
+        prefetch=input_config.get("prefetch", 2),  # Default 2 for bwd compat.
+        **config_to_kw(input_config)
+    )
+    return train_ds, train_data.total_examples
+
+  # A helpful error instead of silent ignore:
+  for k in per_pipeline_configs:
+    assert k not in input_config, f"{k} is per-dataset in multi-input."
+
+  # Parallelize the loading of datasets when doing data mixture.
+  # For larger mixes, we sometimes spend >5min when doing sequentially.
+  # NOTE: functools.cache is thread-safe.
+  def _make(name_and_weight):
+    name, weight = name_and_weight
+    dataset = input_config[name]
+    train_data = ds_core.get(**dataset.data)
+    dataset = make_for_train(
+        data=train_data.get_tfdata(ordered=False, **dataset.get("tfdata", {})),
+        # Don't batch the data just yet, it will be done after
+        # mixing the different datasets below.
+        batch_size=None,
+        preprocess_fn=pp_builder.get_preprocess_fn(dataset.get("pp"), name),
+        prefetch=0,  # Prefetching each pipeline leads to huge OOMs.
+        **config_to_kw(dataset)
+    )
+    if keys := input_config.get("keep_only"):
+      dataset = dataset.map(lambda d, keys=keys: {k: d[k] for k in keys})
+    return name, dataset, weight, train_data.total_examples
+
+  names, datasets, weights, totals = [], [], [], []
+  pool = multiprocessing.pool.ThreadPool(
+      input_config.get("thread_pool_size", len(input_config.data))
+  )
+  for name, dataset, weight, total in pool.map(
+      # Skip weight=0 datasets as a convenient optimization in sweeps.
+      _make, ((name, w) for name, w in input_config.data.items() if w)):
+    names.append(name)
+    datasets.append(dataset)
+    weights.append(weight)
+    totals.append(total)
+
+  # Normalize the weights such that they sum up to 1.
+  weights = [x / sum(weights) for x in weights]
+
+  logging.info(
+      "NOTE: Total dataset mix size: %d\nContributions:\n%s", sum(totals),
+      "\n".join(f"{ds}: {n} ({w * 100:.2g}%)"
+                for ds, n, w in zip(names, totals, weights))
+  )
+
+  train_ds = tf.data.Dataset.sample_from_datasets(
+      datasets, weights, stop_on_empty_dataset=True)
+  if input_config.get("pack"):
+    train_ds = sequence_packing.pack_dataset(
+        train_ds,
+        input_config["batch_size"] // jax.process_count(),
+        input_config.pack.to_dict())
+
+  train_ds = train_ds.batch(
+      input_config["batch_size"] // jax.process_count(), drop_remainder=True)
+  if (pf := input_config.get("prefetch", 2)):
+    train_ds = train_ds.prefetch(pf)
+
+  return train_ds, sum(totals)
+
+
+# The pipeline below is used for evals in multi-{G,T}PU and multi-host settings.
+# As the total number of examples may not be evenly divisible accross all
+# devices, we use the `infinite tf.data padding` trick, which was suggested by
+# Andreas Steiner and also implemented by him in the clu library:
+# https://github.com/google/CommonLoopUtils/blob/84b777c42dfd3fb6685537138433bfeb5241a006/clu/deterministic_data.py#L304.
+def make_for_inference(
+    data, preprocess_fn, batch_size, num_ex_per_process,
+    cache_raw=False, cache_final=False,
+    num_parallel_calls=DEFAULT_NUM_PARALLEL_CALLS, prefetch=1,
+):
+  """Makes an input pipeline for inference."""
+
+  data = _add_tpu_host_options(data)
+  data = data.cache() if cache_raw else data
+  data = data.map(_add_internal_fields(preprocess_fn),
+                  num_parallel_calls=num_parallel_calls)
+  data = data.concatenate(_get_pad_data(data))
+
+  local_batch_size = batch_size // jax.process_count()
+  # This is just like `batch`, but allows batching elements of different shapes
+  # into a tf.RaggedTensor. Elements of the same fixed shape remain tf.Tensors.
+  # Since we do 'infinite' padding it is safe to drop the remainder.
+  data = data.ragged_batch(batch_size=local_batch_size, drop_remainder=True)
+
+  # We need to make sure that all hosts process all data and exactly the same
+  # number of batches. Below we take max per-host num examples and use it on all
+  # hosts to derive the number of batches.
+  num_batches = math.ceil(max(num_ex_per_process) / local_batch_size)
+  data = data.take(num_batches)
+
+  # Note we cache data after a finite number of batches is taken.
+  data = data.cache() if cache_final else data
+  data = data.repeat()
+  data = data.prefetch(prefetch) if prefetch else data
+  return data, num_batches
+
+
+def _get_pad_data(data):
+  def zeros_like_spec(spec):
+    # For unknown/flexible dimensions (None), just use 0 instead.
+    return tf.zeros([x or 0 for x in spec.shape], spec.dtype)
+
+  zero = jax.tree.map(zeros_like_spec, data.element_spec)
+  return tf.data.Dataset.from_tensors(zero).repeat()
+
+
+def _add_internal_fields(pp_fn):
+  """Wraps pp_fn to add _mask and _id keys."""
+  # Adds internal keys, that we either, in this order of preference:
+  # 1. keep from result of pp_fn,
+  # 2. carry over from raw (not pp_fn'd) example, or
+  # 3. add, if that makes sense.
+  def _pp_fn(example):
+    result = pp_fn(example)
+    # _mask will be False on padded examples (see _get_pad_data).
+    result.setdefault("_mask", example.get("_mask", tf.constant(True)))
+    # Not all data-sources can provide an ID. Only carry-over if it can:
+    if "_id" in example and "_id" not in result:
+      result["_id"] = example["_id"]
+    return result
+  return _pp_fn
+
+
+def _add_tpu_host_options(data):
+  options = tf.data.Options()
+  options.threading.private_threadpool_size = 48
+  options.threading.max_intra_op_parallelism = 1
+
+  # Stop a whole bunch of magic stuff that eats up all RAM:
+  options.experimental_optimization.inject_prefetch = False
+
+  return data.with_options(options)
+
+
+def prefetch_iterator(it, n):
+  """Runs iterator `it` ahead for `n` steps. Adapted from flax."""
+  if not n:
+    yield from it
+    return
+  queue = collections.deque()
+
+  def enqueue(n_steps):  # Enqueues *up to* `n` elements from the iterator.
+    for data in itertools.islice(it, n_steps):
+      # Prefetching will parallelize any processing that happens in a different
+      # thread (like `jax.device_put()`), but it will be of no use for
+      # processing that happens in the same thread.
+      queue.append(data)
+
+  enqueue(n)  # Fill up the buffer.
+  while queue:
+    yield queue.popleft()
+    enqueue(1)
+
+
+def threadstart_iterator(it):
+  """Starts an iterator right away in a background thread."""
+  # We already want to "start" the iterator in order to start the underlying
+  # dataset prefetch mechanisms, so here we get the first element. But we don't
+  # want to lose it from training, so we yield that one afterwards.
+  # (internal link)
+  pool = multiprocessing.pool.ThreadPool(processes=1)
+  first_ex_promise = pool.apply_async(lambda: next(it))
+
+  yield first_ex_promise.get()
+  yield from it
+
+
+def tf_to_numpy(x):
+  """Convert any TF types to numpy."""
+  if isinstance(x, tf.Tensor):
+    if x.dtype != tf.string:  # Dense, non-string tensor? Easy!
+      return x.numpy()
+    else:  # A dense string tensor? Turn into actual strings, not bytes.
+      return np.vectorize(bytes.decode, otypes=[str])(x.numpy())
+
+  # The rest deals with RaggedTensors, for two main reasons:
+  # - For strings, recursively apply the above conversion
+  # - For common cases (eg batch of images), return more reasonable shapes.
+
+  # Replace all None's in the shape by a fixed number, in the (somewhat common)
+  # case that they are marked ragged, but really all have the same shape.
+  real_shape = list(x.shape)
+  for i, s in enumerate(real_shape[1:]):
+    if s is not None: continue
+    rowlens = np.diff(x.nested_row_splits[i])
+    if len(set(rowlens)) == 1:
+      real_shape[i + 1] = rowlens[0]
+
+  if None not in real_shape:
+    return tf_to_numpy(x.flat_values).reshape(real_shape)
+
+  # It's actually ragged, reconstruct the array from the variable length pieces.
+  splits = x.row_splits.numpy()
+  rows = [tf_to_numpy(x.values[splits[i]:splits[i + 1]])
+          for i in range(len(splits) - 1)]
+  return np.fromiter(rows, dtype=object)
+
+
+# Note that the order of global devices for sharding data is important and
+# should be compatible with device order used for models params, state, etc.
+def start_global(
+    data, global_devices, n_prefetch=1, keep_on_cpu=frozenset(), warmup=False):
+  """Starts the global input pipeline."""
+  def maybe_shard(name, x):
+    if name in keep_on_cpu:
+      return tf_to_numpy(x)
+    return u.make_fsarray_from_local_slice(x, global_devices)
+
+  it = iter(data)
+  if warmup:  # actually pre-fill shuffle buffers etc.
+    it = threadstart_iterator(it)
+
+  it = (u.tree_map_with_names(maybe_shard, elem) for elem in it)
+  return prefetch_iterator(it, n_prefetch)
+
+
+##########################################################################
+# The code below is pmap-specific and is deprecated, please switch to jit.
+##########################################################################
+
+
+def shard_and_put(x, shard=True, put=True):
+  x = np.asarray(memoryview(x))  # No-copy conversion: http://(internal link)
+  if shard:
+    x = einops.rearrange(x, "(d l) ... -> d l ...", d=jax.local_device_count())
+  if shard and put:  # Only works for pmap (for now).
+    x = jax.device_put_sharded(list(x), jax.local_devices())
+  return x
+
+
+def start_input_pipeline(data, n_prefetch=1, shard=True):
+  fn = functools.partial(shard_and_put, shard=shard, put=n_prefetch)
+  it = (jax.tree.map(fn, elem) for elem in iter(data))
+  return prefetch_iterator(it, n_prefetch)
+
+
+def start_ragged_input_pipeline(data, n_prefetch=1, shard=True, ragged=None):
+  def maybe_shard_and_put(name, x):
+    return x if name in (ragged or {}) else shard_and_put(x, shard)
+
+  it = (u.tree_map_with_names(maybe_shard_and_put, elem) for elem in iter(data))
+  return prefetch_iterator(it, n_prefetch)

Tipsomaly/model/big_vision/load_siglip.py

@@ -0,0 +1,157 @@
+from big_vision.models.proj.image_text import two_towers as model_mod
+
+# Remember that the each module like pp should be imported under the same name 
+# if some other internal files use big_vision.pp and i use .pp then it run through the 
+# module pp twice
+from big_vision.pp import builder as pp_builder 
+from big_vision.pp import ops_general  
+from big_vision.pp import ops_image  
+from big_vision.pp import ops_text  
+from big_vision.pp.proj.image_text import ops_naflex  
+from big_vision.pp.proj.paligemma import ops  
+import PIL
+
+import jax
+import jax.numpy as jnp
+import ml_collections
+import numpy as np
+
+# images = [PIL.Image.open(fname) for fname in [
+#     'apple-ipod.jpg',
+#     'apple-blank.jpg',
+#     'cold_drink.jpg',
+#     'hot_drink.jpg',
+#     'caffeine.jpg',
+#     'siglip.jpg',
+#     'authors.jpg',
+#     'robosign.jpg',
+#     'cow_beach.jpg',
+#     'cow_beach2.jpg',
+#     'mountain_view.jpg',
+# ]]
+# pp_img = pp_builder.get_preprocess_fn(f'resize({RES})|value_range(-1, 1)')
+# imgs = np.array([pp_img({'image': np.array(image)})['image'] for image in images])
+# print('imgs', imgs.shape)
+
+class InputTransform:
+    def __init__(self, RES):
+        self.transform = pp_builder.get_preprocess_fn(f'resize({RES})|value_range(-1, 1)')
+
+    def __call__(self, img):
+        # print(type(img))
+        if np.array(img).size < 3:
+            raise ValueError("invalid image: fewer than 3 elements")
+        return np.array(self.transform({'image': np.array(img)})['image'])
+
+class TargetTransform:
+    def __init__(self, RES):
+        self.transform = pp_builder.get_preprocess_fn(f'resize({RES})|value_range(0, 1)')
+
+    def __call__(self, img):
+        # print(type(img))
+        # print(np.expand_dims(np.array(img), axis=-1).shape)
+
+        out = np.array(self.transform({'image': np.expand_dims(np.array(img), axis=-1)})['image'])
+        # print(out.shape)
+        return np.squeeze(out, axis=-1)
+
+def create_preprocessors_siglip2(RES):
+    # transform = transforms.Compose([
+    #     Ensure3Channels(),
+    #     transforms.Resize((image_size, image_size)),
+    #     transforms.ToTensor(),
+    #     transforms.Normalize(IMAGE_MEAN, IMAGE_STD),
+    # ])
+
+    # target_transform = transforms.Compose([
+    #     transforms.Resize((image_size, image_size)),
+    #     transforms.ToTensor(),
+    # ])
+    pp_input_img = InputTransform(RES)
+    pp_target_img = TargetTransform(RES)
+    return pp_input_img, pp_target_img
+
+def input_transforms(images, RES):
+    pp_img = pp_builder.get_preprocess_fn(f'resize({RES})|value_range(-1, 1)')
+    imgs = np.array([pp_img({'image': np.array(image)})['image'] for image in images])
+    return imgs 
+
+def target_transforms(images, RES):
+    pp_img = pp_builder.get_preprocess_fn(f'resize({RES})|value_range(0, 1)')
+    imgs = np.array([pp_img({'image': np.array(image)})['image'] for image in images])
+    return imgs
+
+def load(VARIANT, RES, ROOT_PATH='/kaggle/working/cpt/'):
+    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 {ROOT_PATH}/{CKPT} || gsutil cp gs://big_vision/siglip2/{CKPT} {ROOT_PATH}
+    # print(f'{ROOT_PATH}/{CKPT} ', f'gs://big_vision/siglip2/{CKPT} ')
+
+    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
+    ))
+    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'/{ROOT_PATH}/{CKPT}', model_cfg)
+
+    return model, params
+
+
+class SigLIPTokenizer:
+    def __init__(self, SEQLEN):
+        self.pp_txt = pp_builder.get_preprocess_fn(f'lower(key="text")|tok(length={SEQLEN}, model="gemma", bos="no", eos="sticky", key="text")')    
+
+    def __call__(self, texts):
+        """texts: str | list[str] -> np.ndarray[int] (B, L)"""
+        if isinstance(texts, str):
+            texts = [texts]
+        return np.array([self.pp_txt({'text': t})['text'] for t in texts])
+
+
+class SigLIPImageEncoder:
+    def __init__(self, model, params):
+        self.model = model
+        self.params = params
+
+    def __call__(self, imgs):
+        return self.model.apply({'params': self.params}, imgs, None)
+
+
+class SigLIPTextEncoder:
+    def __init__(self, model, params):
+        self.model = model
+        self.params = params
+
+    def __call__(self, text_ids, learnable_prompts=None, learning_method=None):
+        return self.model.apply({'params': self.params}, None, text_ids, learnable_prompts=learnable_prompts, learning_method=learning_method)
+
+def build_siglip_modules(model_version, image_size, SEQLEN=64):
+    model, params = load(model_version, image_size)
+    tok = SigLIPTokenizer(SEQLEN)
+    img_enc = SigLIPImageEncoder(model, params)
+    txt_enc = SigLIPTextEncoder(model, params)
+    return img_enc, txt_enc, tok

Tipsomaly/model/big_vision/optax.py

@@ -0,0 +1,225 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Gradient transformations and other optax utilities."""
+
+import operator
+import big_vision.utils as u
+import jax
+import jax.numpy as jnp
+import optax
+
+
+def find_states(opt_state, cls):
+  leaves = jax.tree.leaves(
+      opt_state, is_leaf=lambda node: isinstance(node, cls))
+  return [leaf for leaf in leaves if isinstance(leaf, cls)]
+
+
+def get_count(opt_state, jittable=False):
+  """Returns `ScaleByScheduleState.count` from `opt_state` as an integer."""
+  counts = [
+      state.count
+      for state in find_states(opt_state, optax.ScaleByScheduleState)
+  ]
+  if jittable:
+    return counts[0]
+  else:
+    counts = {int(c) for c in counts}
+    assert len(counts) == 1, f"Expected exactly 1 ScaleByScheduleState:{counts}"
+    return next(iter(counts))
+
+
+def replace_frozen(schedule, pytree, replacement, log=None):
+  """Replaces values matching frozen params in `pytree` with `replacement`."""
+  if not isinstance(schedule, (list, tuple)):
+    return pytree
+  masks, scheds = _make_mask_trees(pytree, schedule, log=log)
+  frozen_mask, _, _ = _split_frozen(masks, scheds)
+  return jax.tree.map(
+      lambda v, f: replacement if f else v, pytree, frozen_mask)
+
+
+def clip_by_per_example_global_norm(
+    max_norm: float,
+) -> optax.GradientTransformation:
+  """Clips the norm of per-example gradients."""
+
+  def init_fn(params):
+    del params
+    return optax.EmptyState()
+
+  def update_fn(updates, state, params=None):
+    del params
+    grads_flat, grads_treedef = jax.tree_util.tree_flatten(updates)
+    batch_size = grads_flat[0].shape[0]
+    clipped, _ = optax.per_example_global_norm_clip(grads_flat, max_norm)
+    grads_sum = jax.tree_util.tree_unflatten(grads_treedef, clipped)
+    grads_mean = jax.tree_util.tree_map(lambda x: x / batch_size, grads_sum)
+    return grads_mean, state
+
+  return optax.GradientTransformation(init_fn, update_fn)
+
+
+def make(config, params, *, sched_kw):
+  """Returns gradient transform and learning rate functions."""
+
+  # Global schedule. No schedule means frozen.
+  schedule = config.get("schedule", {})
+  if not isinstance(schedule, (tuple, list)):
+    schedule = [(".*", schedule)]
+  masks, scheds = _make_mask_trees(params, schedule, "config.schedule")
+  frozen_mask, masks, scheds = _split_frozen(masks, scheds)
+  not_frozen_mask = jax.tree.map(operator.not_, frozen_mask)
+  def create_schedule(mult=1.0, **kw):
+    assert "base" not in kw, kw
+    return u.create_learning_rate_schedule(base=mult, **kw)
+  schedule_fns = [create_schedule(**sched_kw, **sched) for sched in scheds]
+  schedule_txs = [
+      optax.masked(optax.scale_by_schedule(schedule_fn), mask)
+      for schedule_fn, mask in zip(schedule_fns, masks)
+  ] + [
+      # Removes weight decay updates. Note that weight decay already has an
+      # independent mask (which cannot be combined easily with a second mask),
+      # so instead we multiply updates for frozen params with zero.
+      optax.masked(optax.set_to_zero(), frozen_mask)
+  ]
+
+  # Gradient clipping.
+  if clip_norm := config.get("grad_clip_norm"):
+    if config.get("grad_clip_per_example"):
+      clip_tx = clip_by_per_example_global_norm(clip_norm)
+    else:
+      clip_tx = optax.clip_by_global_norm(clip_norm)
+    grad_clip_norm_tx = optax.masked(clip_tx, not_frozen_mask)
+  else:
+    grad_clip_norm_tx = optax.identity()
+
+  # Optimizer updates.
+  tx_func = operator.attrgetter(config.optax_name)(optax)
+  opt_txs = [optax.masked(tx_func(**config.get("optax", {})), not_frozen_mask)]
+  assert "optim" not in config, "Deprecated option, use config.optax."
+
+  # Learning rate multipliers. Defaults to 1.0.
+  lr_mult_txs = [optax.scale(config.lr)]
+  if config.get("lr_mults"):
+    masks, mults = _make_mask_trees(params, config.lr_mults, "config.lr_mults")
+    assert all(mult > 0 for mult in mults), (
+        f"Use schedule=None for parameter freezing instead of lr_mults={mults}")
+    lr_mult_txs += [
+        optax.masked(optax.scale(mult), mask)
+        for mult, mask in zip(mults, masks)
+    ]
+
+  # Weight decay. Defaults to 0.0.
+  # Weight decay is not gradient-based but instead uses "params side-input".
+  # Hence, weight decay is additive and independent of previous gradient-based
+  # updates.
+  assert "weight_decay" not in config, "Deprecated option. Use wd and schedule."
+  assert config.get("weight_decay_decouple", True), (
+      "Coupled weight decay not supported anymore.")
+  if config.get("wd"):
+    wd_mults = config.get("wd_mults", [(".*/kernel$", 1.0)])
+    masks, mults = _make_mask_trees(params, wd_mults, "config.wd_mults")
+    weight_decay_txs = [
+        optax.add_decayed_weights(config.wd * mult, mask)
+        for mult, mask in zip(mults, masks)
+    ]
+  else:
+    weight_decay_txs = []
+
+  # Combine gradient updates and learning rate schedules.
+  return optax.chain(
+      grad_clip_norm_tx,
+      *opt_txs,
+      *lr_mult_txs,
+      *weight_decay_txs,
+      *schedule_txs,
+      optax.scale(-1.0)), schedule_fns
+
+
+def _make_mask_trees(params, patterns_values, log):
+  patterns, values = zip(*patterns_values)
+  masks = u.make_mask_trees(params, patterns, log=log)
+  return masks, values
+
+
+def _split_frozen(masks, scheds):
+  """Computes `frozen_mask` and updates `masks` and `scheds`."""
+  # Specifying `None` as a scheduler freezes params.
+  all_false = jax.tree.map(lambda *bools: not any(bools), *masks)
+  not_covered = [k for k, v in u.tree_flatten_with_names(all_false)[0] if v]
+  assert not not_covered, (
+      f"All params must be covered (use `None` for freezing): {not_covered}")
+  frozen_masks = [
+      mask for mask, sched in zip(masks, scheds) if sched is None]
+  frozen_mask = jax.tree.map(
+      lambda *bools: any(bools), *frozen_masks,
+      all_false)  # `all_false` is required when `frozen_masks==[]`.
+  masks, scheds = zip(*(
+      (mask, sched) for mask, sched in zip(masks, scheds) if sched is not None))
+  return frozen_mask, masks, scheds
+
+
+############ Custom BigVision optimizers #######################################
+# Currently there's only one custom optimizer and we don't foresee new ones in
+# the near future, we opt not to create a new optimizer folder/module for just
+# one isolated case. If there will be more optimizers, we can consider moving
+# them into individual files in a subfolder.
+
+
+# A dummy object to allow for foo.bar access syntax, see
+# https://stackoverflow.com/a/19476841/2366315
+optax.big_vision = type("", (), {})()
+
+
+def scale_by_adafactor(min_dim_size_to_factor=32,
+                       decay_rate=0.8, decay_offset=0,
+                       beta2_cap=0.999,
+                       clipping_threshold=None,
+                       momentum=0.9, dtype_momentum=jnp.bfloat16,
+                       eps=1e-30):
+  """The BigVision variant of Adafactor optimizer."""
+
+  def _decay_rate_pow(i, exponent):
+    """Second-order moment decay schedule."""
+    t = jnp.array(i, jnp.float32) + 1.0
+    return jnp.minimum(beta2_cap, 1.0 - t**(-exponent))
+
+  scale_by_rms = optax.scale_by_factored_rms(
+      factored=True,
+      decay_rate=decay_rate,
+      step_offset=decay_offset,
+      min_dim_size_to_factor=min_dim_size_to_factor,
+      epsilon=eps,
+      decay_rate_fn=_decay_rate_pow)
+
+  clip = (optax.clip_by_block_rms(clipping_threshold) if clipping_threshold
+          else optax.identity())
+
+  mom = (optax.ema(momentum, debias=False, accumulator_dtype=dtype_momentum)
+         if momentum else optax.identity())
+
+  return optax.chain(scale_by_rms, clip, mom)
+
+optax.big_vision.scale_by_adafactor = scale_by_adafactor  # pytype: disable=module-attr
+
+
+# A few more aliases we use frequently:
+def momentum_hp(momentum=0.9, dtype=jnp.bfloat16, nesterov=False):
+  """SGD-Momentum with half-precision accumulator."""
+  return optax.trace(decay=momentum, accumulator_dtype=dtype, nesterov=nesterov)
+
+optax.big_vision.momentum_hp = momentum_hp  # pytype: disable=module-attr
+optax.big_vision.sgd = optax.identity  # pytype: disable=module-attr

Tipsomaly/model/big_vision/optax_test.py

@@ -0,0 +1,341 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for optax."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from big_vision import optax as bv_optax
+import chex
+import jax
+import jax.numpy as jnp
+import ml_collections
+import numpy as np
+import optax
+
+
+class OptaxTest(parameterized.TestCase):
+
+  def test_get_count(self):
+    params = jax.tree.map(jnp.array, {"a": 1.})
+    tx = optax.masked(
+        optax.scale_by_schedule(lambda step: step),
+        {"a": True},
+    )
+    opt_state = tx.init(params)
+    self.assertEqual(bv_optax.get_count(opt_state), 0)
+    _, opt_state = tx.update(params, opt_state)
+    self.assertEqual(bv_optax.get_count(opt_state), 1)
+
+  def test_split_frozen(self):
+    params = jax.tree.map(jnp.array, {
+        "Dense_0": {"kernel": 1., "bias": 2.},
+    })  # pyformat: disable
+    sched1 = dict(decay_type="cosine")
+    sched2 = dict(decay_type="linear")
+    schedule = [
+        (".*/kernel", sched1),
+        (".*/bias", sched2),
+    ]
+    masks, scheds = bv_optax._make_mask_trees(params, schedule, log="schedule")
+    frozen_mask, masks, scheds = bv_optax._split_frozen(masks, scheds)
+    chex.assert_trees_all_equal(
+        frozen_mask,
+        {"Dense_0": {"kernel": False, "bias": False}},
+    )  # pyformat: disable
+    chex.assert_trees_all_equal(
+        masks,
+        (
+            {"Dense_0": {"kernel": True, "bias": False}},
+            {"Dense_0": {"kernel": False, "bias": True}},
+        ),
+    )  # pyformat: disable
+    self.assertEqual(scheds, (sched1, sched2))
+    # freeze some
+    schedule = [
+        (".*/bias", None),
+        ("Dense_0/.*", sched1),
+        (".*", None),
+    ]
+    masks, scheds = bv_optax._make_mask_trees(params, schedule, log="schedule")
+    frozen_mask, masks, scheds = bv_optax._split_frozen(masks, scheds)
+    chex.assert_trees_all_equal(
+        frozen_mask,
+        {"Dense_0": {"kernel": False, "bias": True}},
+    )  # pyformat: disable
+    chex.assert_trees_all_equal(
+        masks,
+        ({"Dense_0": {"kernel": True, "bias": False}},),
+    )  # pyformat: disable
+    self.assertEqual(scheds, (sched1,))
+    # does not cover all params - fails
+    schedule = [
+        (".*/kernel", None),
+    ]
+    masks, scheds = bv_optax._make_mask_trees(params, schedule, log="schedule")
+    with self.assertRaisesRegex(AssertionError, "All params must be covered"):
+      _ = bv_optax._split_frozen(masks, scheds)
+
+  def test_replace_frozen(self):
+    params = jax.tree.map(jnp.array, {
+        "Dense_0": {"kernel": 1., "bias": 2.},
+    })  # pyformat: disable
+    schedule = [
+        (".*/kernel", {}),
+        (".*", None),
+    ]
+    chex.assert_trees_all_equal(
+        bv_optax.replace_frozen(schedule, params, 0.),
+        {"Dense_0": {"kernel": 1., "bias": 0.}},
+    )  # pyformat: disable
+
+  def test_make_simple(self):
+    params = jax.tree.map(jnp.array, {
+        "Dense_0": {"kernel": 1., "bias": 2.},
+    })  # pyformat: disable
+
+    config = ml_collections.ConfigDict()
+    config.lr = 0.01
+    config.schedule = dict(decay_type="linear")
+    config.optax_name = "scale"
+    config.optax = ml_collections.ConfigDict()
+    g_scale = 0.5
+    config.optax.step_size = g_scale
+
+    total_steps = 10
+    sched_kw = dict(global_batch_size=1, total_steps=total_steps)
+    tx, (schedule_fn,) = bv_optax.make(config, params, sched_kw=sched_kw)
+    opt_state = tx.init(params)
+    grads = jax.tree.map(jnp.ones_like, params)
+    for step in range(total_steps):
+      updates, opt_state = tx.update(grads, opt_state)
+      self.assertEqual(bv_optax.get_count(opt_state), step + 1)
+      sched = schedule_fn(step)
+      np.testing.assert_almost_equal(
+          sched, 1.0 / total_steps * (total_steps - step))
+      make_tx = lambda sched: lambda g: -sched * config.lr * g_scale * g
+      chex.assert_trees_all_close(updates, jax.tree.map(make_tx(sched), grads))
+
+  def test_make_wd(self):
+    params = jax.tree.map(jnp.array, {
+        "Dense_0": {"kernel": 1., "bias": 2., "other": 3.},
+    })  # pyformat: disable
+    wds = jax.tree.map(jnp.array, {
+        "Dense_0": {"kernel": 2e-3, "bias": 5e-4, "other": 0.},
+    })  # pyformat: disable
+
+    config = ml_collections.ConfigDict()
+    config.lr = 0.01
+    config.wd = 1e-3
+    config.wd_mults = [
+        (".*/kernel", 2.0),
+        (".*/bias", 0.5),
+    ]
+    config.schedule = dict(decay_type="linear")
+    config.optax_name = "scale"
+    config.optax = ml_collections.ConfigDict()
+    g_scale = 0.5
+    config.optax.step_size = g_scale
+
+    total_steps = 10
+    sched_kw = dict(global_batch_size=1, total_steps=total_steps)
+    tx, (sched_fn,) = bv_optax.make(config, params, sched_kw=sched_kw)
+    opt_state = tx.init(params)
+    grads = jax.tree.map(jnp.ones_like, params)
+    for step in range(total_steps):
+      updates, opt_state = tx.update(grads, opt_state, params)
+      self.assertEqual(bv_optax.get_count(opt_state), step + 1)
+      sched = sched_fn(step)
+      np.testing.assert_almost_equal(
+          sched, 1.0 / total_steps * (total_steps - step))
+
+      def make_tx(sched):
+        def inner(p, g, wd):
+          return -sched * (config.lr * g_scale * g + p * wd)
+        return inner
+
+      chex.assert_trees_all_close(
+          updates, jax.tree.map(make_tx(sched), params, grads, wds))
+
+  def test_make_clip_norm(self):
+    params = jax.tree.map(jnp.array, {
+        "Dense_0": {"kernel": 1., "bias": 2., "other": 3.},
+    })  # pyformat: disable
+
+    config = ml_collections.ConfigDict()
+    config.lr = 0.01
+    config.schedule = dict(decay_type="linear")
+    config.optax_name = "scale"
+    config.grad_clip_norm = 1.0
+    config.optax = ml_collections.ConfigDict()
+    g_scale = 0.5
+    config.optax.step_size = g_scale
+
+    total_steps = 10
+    sched_kw = dict(global_batch_size=1, total_steps=total_steps)
+    tx, (sched_fn,) = bv_optax.make(config, params, sched_kw=sched_kw)
+    opt_state = tx.init(params)
+
+    grads = jax.tree.map(jnp.ones_like, params)
+    gflat = jax.tree.leaves(grads)
+    l2_g = jnp.sqrt(sum([jnp.vdot(p, p) for p in gflat]))
+    grad_clip_factor = jnp.minimum(1.0, config.grad_clip_norm / l2_g)
+    grads_scaled = jax.tree.map(lambda p: grad_clip_factor * p, grads)
+
+    for step in range(total_steps):
+      updates, opt_state = tx.update(grads, opt_state)
+      self.assertEqual(bv_optax.get_count(opt_state), step + 1)
+      sched = sched_fn(step)
+      np.testing.assert_almost_equal(
+          sched, 1.0 / total_steps * (total_steps - step))
+      make_tx = lambda sched: lambda g: -sched * config.lr * g_scale * g
+      chex.assert_trees_all_close(updates,
+                                  jax.tree.map(make_tx(sched), grads_scaled))
+
+  def test_make_multi(self):
+    params = jax.tree.map(
+        jnp.array, {
+            "Dense_0": {"kernel": 1.0, "bias": 2.0, "other": 3.0},
+            "Dense_1": {"kernel": 4.0, "bias": 5.0, "other": 6.0},
+            "Dense_2": {"kernel": 7.0, "bias": 8.0, "other": 9.0},
+            "Dense_3": {"kernel": 10., "bias": 11., "other": 12.},
+        })  # pyformat: disable
+
+    # Manually specify lr + wd for computing expected values.
+    lrb = 0.01
+    lr1 = 2.0
+    lr2 = 0.5
+    lr_mults = {
+        "Dense_0": {"kernel": lr1, "bias": lr1, "other": lr1},
+        "Dense_1": {"kernel": lr2, "bias": lr2, "other": lr2},
+        "Dense_2": {"kernel": 1.0, "bias": 1.0, "other": 1.0},
+        "Dense_3": {"kernel": 1.0, "bias": 1.0, "other": 1.0},
+    }  # pyformat: disable
+    wdb = 1e-3
+    wd1 = 10.0
+    wd2 = 0.1
+    wds = jax.tree.map(
+        jnp.array, {
+            "Dense_0": {"kernel": wd1 * wdb, "bias": wd2 * wdb, "other": 0.},
+            "Dense_1": {"kernel": wd1 * wdb, "bias": wd2 * wdb, "other": 0.},
+            "Dense_2": {"kernel": wd1 * wdb, "bias": wd2 * wdb, "other": 0.},
+            "Dense_3": {"kernel": 0.0 * wdb, "bias": 0.0 * wdb, "other": 0.},
+        })  # pyformat: disable
+
+    config = ml_collections.ConfigDict()
+    config.lr = lrb
+    config.lr_mults = [
+        ("Dense_0/.*", lr1),
+        ("Dense_1/.*", lr2),
+    ]
+    config.wd = wdb
+    config.wd_mults = [
+        (".*/kernel", wd1),
+        (".*/bias", wd2),
+    ]
+    mult1 = 1.0
+    mult2 = 0.1
+    config.schedule = [
+        ("Dense_0/.*", dict(decay_type="linear", mult=mult1, linear_end=mult1)),
+        ("Dense_[12]/.*", dict(decay_type="linear", mult=mult2)),
+        (".*", None),
+    ]
+    config.optax_name = "scale"
+    config.grad_clip_norm = 1.0
+    config.optax = ml_collections.ConfigDict()
+    g_scale = 0.5
+    config.optax.step_size = g_scale
+
+    total_steps = 10
+    sched_kw = dict(global_batch_size=1, total_steps=total_steps)
+    tx, (sched_fn1,
+         sched_fn2) = bv_optax.make(config, params, sched_kw=sched_kw)
+    opt_state = tx.init(params)
+
+    # Manually specify schedules for computing expected values.
+    frozen_fn = lambda _: jnp.array(0.)
+    sched_fns = {
+        "Dense_0": {"kernel": sched_fn1, "bias": sched_fn1, "other": sched_fn1},
+        "Dense_1": {"kernel": sched_fn2, "bias": sched_fn2, "other": sched_fn2},
+        "Dense_2": {"kernel": sched_fn2, "bias": sched_fn2, "other": sched_fn2},
+        "Dense_3": {"kernel": frozen_fn, "bias": frozen_fn, "other": frozen_fn},
+    }  # pyformat: disable
+
+    grads = jax.tree.map(jnp.ones_like, params)
+    gflat, _ = jax.tree.flatten(
+        # Don't count frozen params towards gradient norm.
+        jax.tree.map(lambda g, sched_fn: {frozen_fn: 0}.get(sched_fn, g),
+                     grads, sched_fns))
+    l2_g = jnp.sqrt(sum([jnp.vdot(p, p) for p in gflat]))
+    grad_clip_factor = jnp.minimum(1.0, config.grad_clip_norm / l2_g)
+    grads_scaled = jax.tree.map(lambda p: grad_clip_factor * p, grads)
+
+    def make_tx(step):
+      def get_update(p, g, wd, sched_fn, lr_mult):
+        return -sched_fn(step) * (lrb * lr_mult * g_scale * g + p * wd)
+      return get_update
+
+    for step in range(total_steps):
+      updates, opt_state = tx.update(grads, opt_state, params)
+      self.assertEqual(bv_optax.get_count(opt_state), step + 1)
+      sched1, sched2 = sched_fn1(step), sched_fn2(step)
+      np.testing.assert_almost_equal(sched1, mult1)
+      np.testing.assert_almost_equal(sched2,
+                                     mult2 * (total_steps - step) / total_steps)
+      chex.assert_trees_all_close(
+          updates,
+          jax.tree.map(
+              make_tx(step), params, grads_scaled, wds, sched_fns, lr_mults))
+
+  def test_frozen_no_state(self):
+    params = {"small": jnp.zeros([1]), "large": jnp.zeros([1000])}
+    config = ml_collections.ConfigDict()
+    config.lr = 0.01
+    config.schedule = [
+        ("small", dict(decay_type="cosine")),
+        ("large", None),
+    ]
+    config.optax_name = "scale_by_adam"
+
+    sched_kw = dict(global_batch_size=1, total_steps=1)
+    tx, _ = bv_optax.make(config, params, sched_kw=sched_kw)
+
+    opt_state = tx.init(params)
+    adam_state = bv_optax.find_states(opt_state, optax.ScaleByAdamState)
+    nbytes = sum(
+        jax.tree.flatten(jax.tree.map(lambda x: x.nbytes, adam_state))[0])
+    self.assertLess(nbytes, 1_000)
+
+  def test_adafactor(self):
+    params = {"Dense_0": {"kernel": jnp.zeros([1024, 1024])}}
+
+    config = ml_collections.ConfigDict()
+    config.optax_name = "big_vision.scale_by_adafactor"
+    config.lr = 0.01
+    config.schedule = dict(decay_type="linear")
+    sched_kw = dict(global_batch_size=1, total_steps=1)
+
+    tx, _ = bv_optax.make(config, params, sched_kw=sched_kw)
+
+    opt_state = tx.init(params)
+    adafactor_state = bv_optax.find_states(opt_state, optax.FactoredState)
+    n_state_params = sum(
+        jax.tree.flatten(
+            jax.tree.map(lambda x: np.prod(
+                x.shape if hasattr(x, "shape") else 0), adafactor_state))[0])
+    self.assertEqual(n_state_params, 2 * 1024 + 2)
+
+
+if __name__ == "__main__":
+  absltest.main()

Tipsomaly/model/big_vision/requirements.txt

@@ -0,0 +1,19 @@
+numpy>=1.26
+absl-py
+git+https://github.com/google/CommonLoopUtils
+distrax
+editdistance
+einops
+flax
+optax
+git+https://github.com/google/flaxformer
+git+https://github.com/akolesnikoff/panopticapi.git@mute
+overrides
+protobuf
+sentencepiece
+tensorflow-cpu
+tfds-nightly
+tensorflow-text
+tensorflow-gan
+psutil
+pycocoevalcap

Tipsomaly/model/big_vision/run_tpu.sh

@@ -0,0 +1,35 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#!/bin/bash
+
+if [ ! -d "bv_venv" ]
+then
+  sudo apt-get update
+  sudo apt install -y python3-venv
+  python3 -m venv bv_venv
+  . bv_venv/bin/activate
+
+  pip install -U pip  # Yes, really needed.
+  # NOTE: doesn't work when in requirements.txt -> cyclic dep
+  pip install "jax[tpu]>=0.4.25" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+  pip install -r big_vision/requirements.txt
+else
+  . bv_venv/bin/activate
+fi
+
+if [ $# -ne 0 ]
+then
+  env TFDS_DATA_DIR=$TFDS_DATA_DIR BV_JAX_INIT=1 python3 -m "$@"
+fi

Tipsomaly/model/big_vision/sharding.py

@@ -0,0 +1,197 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Big vision sharding utilities."""
+
+from absl import logging
+
+from big_vision.pp.registry import Registry
+import big_vision.utils as u
+import flax.linen as nn
+import jax
+import numpy as np
+
+
+NamedSharding = jax.sharding.NamedSharding
+P = jax.sharding.PartitionSpec
+
+
+def _replicated(mesh):
+  return NamedSharding(mesh, P())
+
+
+def _shard_along_axis(mesh, i, axis_name):
+  return NamedSharding(mesh, P(*((None,) * i + (axis_name,))))
+
+
+def infer_sharding(params, strategy, mesh):
+  """Infers `params` sharding based on strategy.
+
+  Args:
+    params: a pytree of arrays.
+    strategy: sharding strategy.
+    mesh: jax device mesh.
+
+  Returns:
+    A pytree with shardings, that has the same shape as the `tree` argument.
+  """
+  patterns, tactics = zip(*strategy)
+
+  x_with_names, tree_def = u.tree_flatten_with_names(params)
+  names = tree_def.unflatten(list(zip(*x_with_names))[0])
+
+  # Follows big_vision conventions: each variable is matched at most once,
+  # early patterns get matching priority.
+  mask_trees = u.make_mask_trees(params, patterns)
+
+  specs = jax.tree.map(lambda x: (None,) * x.ndim, params)
+
+  for mask_tree, tactic in zip(mask_trees, tactics):
+    for op_str in tactic.split("|"):
+      op = Registry.lookup(f"shardings.{op_str}")()
+      specs = jax.tree.map(
+          lambda x, n, match, spec, op=op: op(spec, mesh, n, x)
+          if match else spec,
+          params, names, mask_tree, specs,
+          is_leaf=lambda v: isinstance(v, nn.Partitioned))
+
+  # Two-level tree_map to prevent it from doing traversal inside the spec.
+  specs = jax.tree.map(lambda _, spec: P(*spec), nn.unbox(params), specs)
+  return jax.tree.map(lambda spec: NamedSharding(mesh, spec), specs)
+
+
+# Sharding rules
+#
+# Each rule needs to be added to the registry, can accept custom args, and
+# returns a function that updates the current spec. The arguments are:
+# 1. Variable name
+# 2. Variable itself (or placeholder with .shape and .dtype properties)
+# 3. The current sharing spec.
+
+
+@Registry.register("shardings.replicate")
+def replicate():
+  """Full replication sharding rule.
+
+  Note full replication is deafult, so this can be skipped and useful to
+  explicitly state in the config that certrain parameters are replicated.
+  TODO: can be generalized to support replication over a sub-mesh.
+
+  Returns:
+    A function that updates the sharding spec.
+  """
+  def _update_spec(cur_spec, mesh, name, x):
+    del x, mesh
+    if not all(axis is None for axis in cur_spec):
+      raise ValueError(f"Inconsistent sharding instructions: "
+                       f"parameter {name} has spec {cur_spec}, "
+                       f"so it can't be fully replicated.")
+    return cur_spec
+  return _update_spec
+
+
+@Registry.register("shardings.fsdp")
+def fsdp(axis, min_size_to_shard_mb=4):
+  """FSDP sharding rule.
+
+  Shards the largest dimension that is not sharded already and is divisible
+  by the total device count.
+
+  Args:
+    axis: mesh axis name for FSDP, or a collection of names.
+    min_size_to_shard_mb: minimal tensor size to bother with sharding.
+
+  Returns:
+    A function that updates the sharding spec.
+  """
+  axis = axis if isinstance(axis, str) else tuple(axis)
+  axis_tuple = axis if isinstance(axis, tuple) else (axis,)
+  def _update_spec(cur_spec, mesh, name, x):
+    shape = x.shape
+    axis_size = np.prod([mesh.shape[a] for a in axis_tuple])
+
+    if np.prod(shape) * x.dtype.itemsize <= min_size_to_shard_mb * (2 ** 20):
+      return cur_spec
+
+    # Partition along largest axis that is divisible and not taken.
+    idx = np.argsort(shape)[::-1]
+    for i in idx:
+      if shape[i] % axis_size == 0:
+        if cur_spec[i] is None:
+          return cur_spec[:i] + (axis,) + cur_spec[i+1:]
+
+    logging.info("Failed to apply `fsdp` rule to the parameter %s:%s, as all "
+                 "its dimensions are not divisible by the requested axis: "
+                 "%s:%i, or already occupied by other sharding rules: %s",
+                 name, shape, axis, axis_size, cur_spec)
+    return cur_spec
+  return _update_spec
+
+
+@Registry.register("shardings.logical_partitioning")
+def logical_partitioning():
+  """Manual sharding based on Flax's logical partitioning annotations.
+
+  Uses logical sharding annotations added in model code with
+  `nn.with_logical_partitioning`.  Respects logical to mesh name mapping rules
+  (typically defined in the dynamic context using
+  `with nn.logical_axis_rules(rules): ...`).
+
+  Returns:
+    A function that outputs the sharding spec of `nn.LogicallyPartitioned` boxed
+    specs.
+  """
+  def _update_spec(cur_spec, mesh, name, x):
+    del x, name, mesh
+    if isinstance(cur_spec, nn.LogicallyPartitioned):
+      return nn.logical_to_mesh_axes(cur_spec.names)
+    return cur_spec
+  return _update_spec
+
+
+@Registry.register("shardings.shard_dim")
+def shard_dim(axis, dim, ignore_ndim_error=False):
+  """Shards the given dimension along the given axis.
+
+  Args:
+    axis: mesh axis name for sharding.
+    dim: dimension to shard (can be negative).
+    ignore_ndim_error: if True, a warning error is logged instead of raising an
+      exception when the given dimension is not compatible with the number of
+      dimensions of the array.
+
+  Returns:
+    A function that updates the sharding spec.
+  """
+  def _update_spec(cur_spec, mesh, name, x):
+    del mesh, x
+    if np.abs(dim) >= len(cur_spec):
+      msg = f"Cannot shard_dim({axis}, {dim}): name={name} cur_spec={cur_spec}"
+      if ignore_ndim_error:
+        logging.warning(msg)
+        return cur_spec
+      else:
+        raise ValueError(msg)
+    pos_dim = dim
+    if pos_dim < 0:
+      pos_dim += len(cur_spec)
+    if cur_spec[pos_dim] is not None:
+      raise ValueError(
+          f"Already sharded: shard_dim({axis}, {dim}):"
+          f" name={name} cur_spec={cur_spec}"
+      )
+    new_spec = cur_spec[:pos_dim] + (axis,) + cur_spec[pos_dim + 1 :]
+    return new_spec
+
+  return _update_spec

Tipsomaly/model/big_vision/train.py

@@ -0,0 +1,518 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Training loop example.
+
+This is a basic variant of a training loop, good starting point for fancy ones.
+"""
+# pylint: disable=consider-using-from-import
+# pylint: disable=logging-fstring-interpolation
+
+import functools
+import importlib
+import multiprocessing.pool
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+import big_vision.evaluators.common as eval_common
+import big_vision.input_pipeline as input_pipeline
+import big_vision.optax as bv_optax
+import big_vision.sharding as bv_sharding
+import big_vision.utils as u
+from clu import parameter_overview
+import flax.linen as nn
+import jax
+from jax.experimental import multihost_utils
+from jax.experimental.array_serialization import serialization as array_serial
+from jax.experimental.shard_map import shard_map
+import jax.numpy as jnp
+from ml_collections import config_flags
+import numpy as np
+import optax
+import tensorflow as tf
+
+from tensorflow.io import gfile
+
+
+config_flags.DEFINE_config_file(
+    "config", None, "Training configuration.", lock_config=True)
+
+flags.DEFINE_string("workdir", default=None, help="Work unit directory.")
+flags.DEFINE_boolean("cleanup", default=False,
+                     help="Delete workdir (only) after successful completion.")
+
+# Adds jax flags to the program.
+jax.config.parse_flags_with_absl()
+# Transfer guard will fail the program whenever that data between a host and
+# a device is transferred implicitly. This often catches subtle bugs that
+# cause slowdowns and memory fragmentation. Explicit transfers are done
+# with jax.device_put and jax.device_get.
+jax.config.update("jax_transfer_guard", "disallow")
+# Fixes design flaw in jax.random that may cause unnecessary d2d comms.
+jax.config.update("jax_threefry_partitionable", True)
+
+
+NamedSharding = jax.sharding.NamedSharding
+P = jax.sharding.PartitionSpec
+
+
+def main(argv):
+  del argv
+
+  # This is needed on multihost systems, but crashes on non-TPU single-host.
+  if os.environ.get("BV_JAX_INIT"):
+    jax.distributed.initialize()
+
+  # Make sure TF does not touch GPUs.
+  tf.config.set_visible_devices([], "GPU")
+
+  config = flags.FLAGS.config
+
+################################################################################
+#                                                                              #
+#                                Set up logging                                #
+#                                                                              #
+################################################################################
+
+  # Set up work directory and print welcome message.
+  workdir = flags.FLAGS.workdir
+  logging.info(
+      f"\u001b[33mHello from process {jax.process_index()} holding "
+      f"{jax.local_device_count()}/{jax.device_count()} devices and "
+      f"writing to workdir {workdir}.\u001b[0m")
+  logging.info(f"The config:\n{config}")
+
+  save_ckpt_path = None
+  if workdir:  # Always create if requested, even if we may not write into it.
+    gfile.makedirs(workdir)
+    save_ckpt_path = os.path.join(workdir, "checkpoint.bv")
+
+  # The pool is used to perform misc operations such as logging in async way.
+  pool = multiprocessing.pool.ThreadPool(1)
+
+  # Here we register preprocessing ops from modules listed on `pp_modules`.
+  for m in config.get("pp_modules", ["ops_general", "ops_image", "ops_text"]):
+    importlib.import_module(f"big_vision.pp.{m}")
+
+  # Setup up logging and experiment manager.
+  xid, wid = -1, -1
+  fillin = lambda s: s
+  def info(s, *a):
+    logging.info("\u001b[33mNOTE\u001b[0m: " + s, *a)
+  def write_note(note):
+    if jax.process_index() == 0:
+      info("%s", note)
+
+  mw = u.BigVisionMetricWriter(xid, wid, workdir, config)
+
+  # Allow for things like timings as early as possible!
+  u.chrono.inform(measure=mw.measure, write_note=write_note)
+
+################################################################################
+#                                                                              #
+#                                Set up Mesh                                   #
+#                                                                              #
+################################################################################
+
+  # We rely on jax mesh_utils to organize devices, such that communication
+  # speed is the fastest for the last dimension, second fastest for the
+  # penultimate dimension, etc.
+  config_mesh = config.get("mesh", [("data", jax.device_count())])
+
+  # Sharding rules with default
+  sharding_rules = config.get("sharding_rules", [("act_batch", "data")])
+
+  write_note("Creating device mesh...")
+  mesh = u.create_device_mesh(
+      config_mesh,
+      allow_split_physical_axes=config.get("mesh_allow_split_physical_axes",
+                                           False))
+  repl_sharding = jax.sharding.NamedSharding(mesh, P())
+
+  # Consistent device order is important to ensure correctness of various train
+  # loop components, such as input pipeline, update step, evaluators. The
+  # order presribed by the `devices_flat` variable should be used throughout
+  # the program.
+  devices_flat = mesh.devices.flatten()
+
+################################################################################
+#                                                                              #
+#                                Input Pipeline                                #
+#                                                                              #
+################################################################################
+
+  write_note("Initializing train dataset...")
+  batch_size = config.input.batch_size
+  if batch_size % jax.device_count() != 0:
+    raise ValueError(f"Batch size ({batch_size}) must "
+                     f"be divisible by device number ({jax.device_count()})")
+  info("Global batch size %d on %d hosts results in %d local batch size. With "
+       "%d dev per host (%d dev total), that's a %d per-device batch size.",
+       batch_size, jax.process_count(), batch_size // jax.process_count(),
+       jax.local_device_count(), jax.device_count(),
+       batch_size // jax.device_count())
+
+  train_ds, ntrain_img = input_pipeline.training(config.input)
+
+  total_steps = u.steps("total", config, ntrain_img, batch_size)
+  def get_steps(name, default=ValueError, cfg=config):
+    return u.steps(name, cfg, ntrain_img, batch_size, total_steps, default)
+
+  u.chrono.inform(total_steps=total_steps, global_bs=batch_size,
+                  steps_per_epoch=ntrain_img / batch_size)
+
+  info("Running for %d steps, that means %f epochs",
+       total_steps, total_steps * batch_size / ntrain_img)
+
+  # Start input pipeline as early as possible.
+  n_prefetch = config.get("prefetch_to_device", 1)
+  train_iter = input_pipeline.start_global(train_ds, devices_flat, n_prefetch)
+
+################################################################################
+#                                                                              #
+#                           Create Model & Optimizer                           #
+#                                                                              #
+################################################################################
+
+  write_note("Creating model...")
+  model_mod = importlib.import_module(f"big_vision.models.{config.model_name}")
+  model = model_mod.Model(
+      num_classes=config.num_classes, **config.get("model", {}))
+
+  def init(rng):
+    batch = jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype.as_numpy_dtype),
+                         train_ds.element_spec)
+    params = model.init(rng, batch["image"])["params"]
+
+    # Set bias in the head to a low value, such that loss is small initially.
+    if "init_head_bias" in config:
+      params["head"]["bias"] = jnp.full_like(params["head"]["bias"],
+                                             config["init_head_bias"])
+
+    return params
+
+  # This seed makes the Jax part of things (like model init) deterministic.
+  # However, full training still won't be deterministic, for example due to the
+  # tf.data pipeline not being deterministic even if we would set TF seed.
+  # See (internal link) for a fun read on what it takes.
+  rng = jax.random.PRNGKey(u.put_cpu(config.get("seed", 0)))
+
+  write_note("Inferring parameter shapes...")
+  rng, rng_init = jax.random.split(rng)
+  params_shape = jax.eval_shape(init, rng_init)
+
+  write_note("Inferring optimizer state shapes...")
+  tx, sched_fns = bv_optax.make(config, nn.unbox(params_shape), sched_kw=dict(
+      total_steps=total_steps, batch_size=batch_size, data_size=ntrain_img))
+  opt_shape = jax.eval_shape(tx.init, params_shape)
+  # We jit this, such that the arrays are created on the CPU, not device[0].
+  sched_fns_cpu = [u.jit_cpu()(sched_fn) for sched_fn in sched_fns]
+
+  if jax.process_index() == 0:
+    num_params = sum(np.prod(p.shape) for p in jax.tree.leaves(params_shape))
+    mw.measure("num_params", num_params)
+
+################################################################################
+#                                                                              #
+#                               Shard & Transfer                               #
+#                                                                              #
+################################################################################
+
+  write_note("Inferring shardings...")
+  train_state_shape = {"params": params_shape, "opt": opt_shape}
+
+  strategy = config.get("sharding_strategy", [(".*", "replicate")])
+  with nn.logical_axis_rules(sharding_rules):
+    train_state_sharding = bv_sharding.infer_sharding(
+        train_state_shape, strategy=strategy, mesh=mesh)
+
+  write_note("Transferring train_state to devices...")
+  # RNG is always replicated
+  rng_init = u.reshard(rng_init, repl_sharding)
+
+  # Parameters and the optimizer are now global (distributed) jax arrays.
+  params = jax.jit(init, out_shardings=train_state_sharding["params"])(rng_init)
+  opt = jax.jit(tx.init, out_shardings=train_state_sharding["opt"])(params)
+
+  rng, rng_loop = jax.random.split(rng, 2)
+  rng_loop = u.reshard(rng_loop, repl_sharding)
+  del rng  # not used anymore, so delete it.
+
+  # At this point we have everything we need to form a train state. It contains
+  # all the parameters that are passed and updated by the main training step.
+  # From here on, we have no need for Flax AxisMetadata (such as partitioning).
+  train_state = nn.unbox({"params": params, "opt": opt})
+  del params, opt  # Delete to avoid memory leak or accidental reuse.
+
+  write_note("Logging parameter overview...")
+  parameter_overview.log_parameter_overview(
+      train_state["params"], msg="Init params",
+      include_stats="global", jax_logging_process=0)
+
+################################################################################
+#                                                                              #
+#                                 Update Step                                  #
+#                                                                              #
+################################################################################
+
+  @functools.partial(
+      jax.jit,
+      donate_argnums=(0,),
+      out_shardings=(train_state_sharding, repl_sharding))
+  def update_fn(train_state, rng, batch):
+    """Update step."""
+
+    images, labels = batch["image"], batch["labels"]
+
+    step_count = bv_optax.get_count(train_state["opt"], jittable=True)
+    rng = jax.random.fold_in(rng, step_count)
+
+    if config.get("mixup") and config.mixup.p:
+      # The shard_map below makes mixup run on every device independently and
+      # thus avoids unnecessary communication.
+      sharded_mixup_fn = shard_map(
+          u.get_mixup(rng, config.mixup.p),
+          mesh=jax.sharding.Mesh(devices_flat, ("data",)),
+          in_specs=P("data"), out_specs=(P(), P("data"), P("data")))
+      rng, (images, labels), _ = sharded_mixup_fn(images, labels)
+
+    # Get device-specific loss rng.
+    rng, rng_model = jax.random.split(rng, 2)
+
+    def loss_fn(params):
+      logits, _ = model.apply(
+          {"params": params}, images,
+          train=True, rngs={"dropout": rng_model})
+      return getattr(u, config.get("loss", "sigmoid_xent"))(
+          logits=logits, labels=labels)
+
+    params, opt = train_state["params"], train_state["opt"]
+    loss, grads = jax.value_and_grad(loss_fn)(params)
+    updates, opt = tx.update(grads, opt, params)
+    params = optax.apply_updates(params, updates)
+
+    measurements = {"training_loss": loss}
+    gs = jax.tree.leaves(bv_optax.replace_frozen(config.schedule, grads, 0.))
+    measurements["l2_grads"] = jnp.sqrt(sum([jnp.sum(g * g) for g in gs]))
+    ps = jax.tree.leaves(params)
+    measurements["l2_params"] = jnp.sqrt(sum([jnp.sum(p * p) for p in ps]))
+    us = jax.tree.leaves(updates)
+    measurements["l2_updates"] = jnp.sqrt(sum([jnp.sum(u * u) for u in us]))
+
+    return {"params": params, "opt": opt}, measurements
+
+################################################################################
+#                                                                              #
+#                               Load Checkpoint                                #
+#                                                                              #
+################################################################################
+
+  # Decide how to initialize training. The order is important.
+  # 1. Always resumes from the existing checkpoint, e.g. resumes a finetune job.
+  # 2. Resume from a previous checkpoint, e.g. start a cooldown training job.
+  # 3. Initialize model from something, e,g, start a fine-tuning job.
+  # 4. Train from scratch.
+  resume_ckpt_path = None
+  if save_ckpt_path and gfile.exists(f"{save_ckpt_path}-LAST"):
+    resume_ckpt_path = save_ckpt_path
+  elif config.get("resume"):
+    resume_ckpt_path = fillin(config.resume)
+
+  ckpt_mngr = None
+  if save_ckpt_path or resume_ckpt_path:
+    ckpt_mngr = array_serial.GlobalAsyncCheckpointManager()
+
+  if resume_ckpt_path:
+    write_note(f"Resuming training from checkpoint {resume_ckpt_path}...")
+    jax.tree.map(lambda x: x.delete(), train_state)
+    del train_state
+    shardings = {
+        **train_state_sharding,
+        "chrono": jax.tree.map(lambda _: repl_sharding,
+                               u.chrono.save()),
+    }
+    loaded = u.load_checkpoint_ts(
+        resume_ckpt_path, tree=shardings, shardings=shardings)
+    train_state = {key: loaded[key] for key in train_state_sharding.keys()}
+
+    u.chrono.load(jax.device_get(loaded["chrono"]))
+    del loaded
+  elif config.get("model_init"):
+    write_note(f"Initialize model from {config.model_init}...")
+    # TODO: when updating the `load` API soon, do pass and request the
+    # full `train_state` from it. Examples where useful: VQVAE, BN.
+    train_state["params"] = model_mod.load(
+        train_state["params"], config.model_init, config.get("model"),
+        **config.get("model_load", {}))
+
+    # load has the freedom to return params not correctly sharded. Think of for
+    # example ViT resampling position embedings on CPU as numpy arrays.
+    train_state["params"] = u.reshard(
+        train_state["params"], train_state_sharding["params"])
+
+    parameter_overview.log_parameter_overview(
+        train_state["params"], msg="restored params",
+        include_stats="global", jax_logging_process=0)
+
+
+################################################################################
+#                                                                              #
+#                                 Setup Evals                                  #
+#                                                                              #
+################################################################################
+
+  # We do not jit/pmap this function, because it is passed to evaluator that
+  # does it later. We output as many intermediate tensors as possible for
+  # maximal flexibility. Later `jit` will prune out things that are not needed.
+  def eval_logits_fn(train_state, batch):
+    logits, out = model.apply({"params": train_state["params"]}, batch["image"])
+    return logits, out
+
+  def eval_loss_fn(train_state, batch):
+    logits, _ = model.apply({"params": train_state["params"]}, batch["image"])
+    loss_fn = getattr(u, config.get("loss", "sigmoid_xent"))
+    return {
+        "loss": loss_fn(logits=logits, labels=batch["labels"], reduction=False)
+    }
+
+  eval_fns = {
+      "predict": eval_logits_fn,
+      "loss": eval_loss_fn,
+  }
+
+  # Only initialize evaluators when they are first needed.
+  @functools.lru_cache(maxsize=None)
+  def evaluators():
+    return eval_common.from_config(
+        config, eval_fns,
+        lambda s: write_note(f"Init evaluator: {s}…\n{u.chrono.note}"),
+        lambda key, cfg: get_steps(key, default=None, cfg=cfg),
+        devices_flat,
+    )
+
+  # At this point we need to know the current step to see whether to run evals.
+  write_note("Inferring the first step number...")
+  first_step_device = bv_optax.get_count(train_state["opt"], jittable=True)
+  first_step = int(jax.device_get(first_step_device))
+  u.chrono.inform(first_step=first_step)
+
+  # Note that training can be pre-empted during the final evaluation (i.e.
+  # just after the final checkpoint has been written to disc), in which case we
+  # want to run the evals.
+  if first_step in (total_steps, 0):
+    write_note("Running initial or final evals...")
+    mw.step_start(first_step)
+    for (name, evaluator, _, prefix) in evaluators():
+      if config.evals[name].get("skip_first") and first_step != total_steps:
+        continue
+      write_note(f"{name} evaluation...\n{u.chrono.note}")
+      with u.chrono.log_timing(f"z/secs/eval/{name}"):
+        with mesh, nn.logical_axis_rules(sharding_rules):
+          for key, value in evaluator.run(train_state):
+            mw.measure(f"{prefix}{key}", value)
+
+################################################################################
+#                                                                              #
+#                                  Train Loop                                  #
+#                                                                              #
+################################################################################
+
+  prof = None  # Keeps track of start/stop of profiler state.
+
+  write_note("Starting training loop, compiling the first step...")
+  for step, batch in zip(range(first_step + 1, total_steps + 1), train_iter):
+    mw.step_start(step)
+
+    with jax.profiler.StepTraceAnnotation("train_step", step_num=step):
+      with u.chrono.log_timing("z/secs/update0", noop=step > first_step + 1):
+        with mesh, nn.logical_axis_rules(sharding_rules):
+          train_state, measurements = update_fn(train_state, rng_loop, batch)
+
+    # On the first host, let's always profile a handful of early steps.
+    if jax.process_index() == 0:
+      prof = u.startstop_prof(prof, step, first_step, get_steps("log_training"))
+
+    # Report training progress
+    if (u.itstime(step, get_steps("log_training"), total_steps, host=0)
+        or u.chrono.warmup and jax.process_index() == 0):
+      for i, sched_fn_cpu in enumerate(sched_fns_cpu):
+        mw.measure(f"global_schedule{i if i else ''}",
+                   sched_fn_cpu(u.put_cpu(step - 1)))
+      measurements = jax.device_get(measurements)
+      for name, value in measurements.items():
+        mw.measure(name, value)
+      u.chrono.tick(step)
+      for k in ("training_loss", "l2_grads", "l2_updates", "l2_params"):
+        if not np.isfinite(measurements.get(k, 0.0)):
+          raise RuntimeError(f"{k} became nan or inf somewhere within steps "
+                             f"[{step - get_steps('log_training')}, {step}]")
+
+    # Checkpoint saving
+    keep_last = total_steps if get_steps("ckpt", None) else None
+    keep_ckpt_steps = get_steps("keep_ckpt", None) or keep_last
+    if save_ckpt_path and (
+        (keep := u.itstime(step, keep_ckpt_steps, total_steps, first=False))
+        or u.itstime(step, get_steps("ckpt", None), total_steps, first=True)
+    ):
+      u.chrono.pause(wait_for=train_state)
+
+      # Copy because we add extra stuff to the checkpoint.
+      ckpt = {**train_state}
+
+      # To save chrono state correctly and safely in a multihost setup, we
+      # broadcast the state to all hosts and convert it to a global array.
+      with jax.transfer_guard("allow"):
+        chrono_ckpt = multihost_utils.broadcast_one_to_all(u.chrono.save())
+      chrono_shardings = jax.tree.map(lambda _: repl_sharding, chrono_ckpt)
+      ckpt = ckpt | {"chrono": u.reshard(chrono_ckpt, chrono_shardings)}
+
+      u.save_checkpoint_ts(ckpt_mngr, ckpt, save_ckpt_path, step, keep)
+      u.chrono.resume()
+
+    for (name, evaluator, log_steps, prefix) in evaluators():
+      if u.itstime(step, log_steps, total_steps, first=False, last=True):
+        u.chrono.pause(wait_for=train_state)
+        u.chrono.tick(step)  # Record things like epoch number, core hours etc.
+        write_note(f"{name} evaluation...\n{u.chrono.note}")
+        with u.chrono.log_timing(f"z/secs/eval/{name}"):
+          with mesh, nn.logical_axis_rules(sharding_rules):
+            for key, value in evaluator.run(train_state):
+              mw.measure(f"{prefix}{key}", jax.device_get(value))
+        u.chrono.resume()
+    mw.step_end()
+
+  # Always give a chance to stop the profiler, no matter how things ended.
+  # TODO: can we also do this when dying of an exception like OOM?
+  if jax.process_index() == 0 and prof is not None:
+    u.startstop_prof(prof)
+
+  # Last note needs to happen before the pool's closed =)
+  write_note(f"Done!\n{u.chrono.note}")
+
+  pool.close()
+  pool.join()
+  mw.close()
+  if ckpt_mngr:
+    ckpt_mngr.wait_until_finished()
+
+  # Make sure all hosts stay up until the end of main.
+  u.sync()
+
+  u.maybe_cleanup_workdir(workdir, flags.FLAGS.cleanup, info)
+
+
+if __name__ == "__main__":
+  app.run(main)

Tipsomaly/model/big_vision/utils.py

@@ -0,0 +1,1478 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utils very specific to this project, not generic."""
+
+import collections
+import contextlib
+import dataclasses
+import functools
+import io
+import json
+import multiprocessing
+import multiprocessing.pool
+import os
+import re
+import sys
+import time
+from typing import Mapping
+
+from absl import flags
+from absl import logging
+from big_vision.pp import registry as pp_registry
+import einops
+import flax
+import flax.jax_utils as flax_utils
+import jax
+from jax.experimental import mesh_utils
+from jax.experimental.array_serialization import serialization as array_serial
+import jax.numpy as jnp
+import ml_collections as mlc
+import numpy as np
+
+import tensorflow.io.gfile as gfile  # pylint: disable=consider-using-from-import
+
+
+Registry = pp_registry.Registry
+
+
+# pylint: disable=logging-fstring-interpolation
+
+
+def pad_shard_unpad(wrapped, static_argnums=(0,), static_argnames=()):
+  """Wraps a function with code that pads, shards, then un-shards, un-pads.
+
+  Args:
+    wrapped: the function to be wrapped. Signature is `params, *args, *kwargs`.
+    static_argnums: indices of arguments to `wrapped` that should _not_ be
+      padded and sharded, but instead be forwarded as-is. The default is (0,)
+      because by far the most common use-case is to pass `params` first.
+    static_argnames: names of kwargs to `wrapped` that should _not_ be padded
+      and sharded, but instead be forwarded as-is.
+
+  Returns:
+    A new function that pads and shards its arguments before passing them to
+    the wrapped function, and un-shards and un-pads the returned pytree.
+
+    This is useful for calling a pmap'ed function with inputs that aren't
+    divisible by the number of devices. A typical use is:
+      @pad_shard_unpad
+      @jax.pmap
+      def forward(params, x): ...
+
+  Notes:
+    The padding is done in host-memory before being passed to the function, and
+    the values returned by the function are transferred back to host memory.
+
+    The returned function is augmented with a new keyword-only argument
+    `min_device_batch` that, if specified, forces padding inputs to at least
+    this size per device. This can be useful to avoid recompiles for the last
+    batch and reduce memory fragmentation.
+  """
+
+  def pad_shard_unpad_wrapper(*args, min_device_batch=None, **kw):
+    d = jax.local_device_count()  # d = devices, b = batch
+
+    # Find the batch-sizes of all non-static arguments.
+    def get_bs(x):
+      batch_sizes = jax.tree.map(lambda y: y.shape[0], x)
+      return jax.tree.flatten(batch_sizes)[0]
+
+    bs_a = [get_bs(a) for i, a in enumerate(args) if i not in static_argnums]
+    bs_kw = [get_bs(v) for k, v in kw.items() if k not in static_argnames]
+    bs = set([n for b in (bs_a + bs_kw) for n in b])
+    assert len(bs) == 1, f"Inconsistent batch-sizes: {bs}"
+    b = bs.pop()
+
+    def pad(x):
+      _, *shape = x.shape
+      db, rest = divmod(b, d)
+      if rest:
+        x = np.concatenate([x, np.zeros((d - rest, *shape), x.dtype)], axis=0)
+        db += 1
+      if min_device_batch and db < min_device_batch:
+        x = np.concatenate(
+            [x, np.zeros((d * (min_device_batch - db), *shape), x.dtype)])
+        db = min_device_batch
+      return x.reshape(d, db, *shape)
+
+    def maybe_pad(x, actually_pad=True):
+      if not actually_pad: return x  # For call-site convenience below.
+      return jax.tree.map(pad, x)
+
+    args = [maybe_pad(a, i not in static_argnums) for i, a in enumerate(args)]
+    kw = {k: maybe_pad(v, k not in static_argnames) for k, v in kw.items()}
+    out = wrapped(*args, **kw)
+
+    def unpad(x):
+      # Transfer back before cutting, to reduce on-device shape diversity.
+      return einops.rearrange(jax.device_get(x), "d b ... -> (d b) ...")[:b]
+    return jax.tree.map(unpad, out)
+
+  return pad_shard_unpad_wrapper
+
+
+def onehot(labels, num_classes, on_value=1.0, off_value=0.0):
+  x = (labels[..., None] == jnp.arange(num_classes)[None])
+  x = jax.lax.select(x, jnp.full(x.shape, on_value),
+                     jnp.full(x.shape, off_value))
+  return x.astype(jnp.float32)
+
+
+def npload(fname):
+  """Loads `fname` and returns an np.ndarray or dict thereof."""
+  # Load the data; use local paths directly if possible:
+  if os.path.exists(fname):
+    loaded = np.load(fname, allow_pickle=False)
+  else:
+    # For other (remote) paths go via gfile+BytesIO as np.load requires seeks.
+    with gfile.GFile(fname, "rb") as f:
+      data = f.read()
+    loaded = np.load(io.BytesIO(data), allow_pickle=False)
+
+  # Support loading both single-array files (np.save) and zips (np.savez).
+  if isinstance(loaded, np.ndarray):
+    return loaded
+  else:
+    return dict(loaded)
+
+
+def load_checkpoint_np(npz, tree=None):
+  """Loads a jax pytree from a npz file.
+
+  Args:
+    npz: Either path to the checkpoint file (.npz), or a dict-like.
+    tree: deprecated, use None.
+      Bwd-compat for old format that only stored values: the pytree structure.
+
+  Returns:
+    A pytree that is the checkpoint.
+  """
+  if isinstance(npz, str):  # If not already loaded, then load.
+    npz = npload(npz)
+  keys, values = zip(*list(npz.items()))
+  if tree:
+    checkpoint = tree.unflatten(values)
+  else:
+    checkpoint = recover_tree(keys, values)
+  return checkpoint
+
+
+def load_params(ckpt, **kw):
+  """Loads the parameters of a big_vision checkpoint, both old or new format.
+
+  Args:
+    ckpt: Path to the checkpoint (.npz, .ts) or dict-like.
+    **kw: forwarded to the underlying load function (_np or _ts).
+
+  Returns:
+    A pytree that is the checkpoint, potentially sharded.
+
+  Notes:
+    The `ckpt` string can contain an colon-separated "submodel" indicator, like
+    `img` in the example `/path/to/file.npz:img`.
+    This is used to load sub-parts of a model, for example the image load the
+    image encoder out of a two_tower (SigLIP) checkpoint, or distillation.
+    This way, ANY model that uses this function can load itself from a
+    checkpoint that contains multiple sub-models.
+  """
+  key = None  # Whether we want to extract only a sub-key of the model.
+
+  if isinstance(ckpt, str):  # Most common case of passing a checkpoint path.
+    # Potentially read out the sub-part to load from after the colon
+    # '/path/to/file:img/head' => '/path/to/file', 'img/head'
+    # 'gs://path/to/file' => 'gs://path/to/file', None
+    if match := re.match(r"^(.*?/.*?)(?::([\w/]+))?$", ckpt):
+      ckpt, key = match.groups()
+    else:
+      raise ValueError(f"Weird ckpt path: {ckpt} ; Maybe prepend ./ ?")
+
+    # Use the checkpoint filename to detect when we're loading old-style .npz
+    # checkpoints, as opposed to new-style tensorstore checkpoint folders.
+    if ".npz" in ckpt:  # Not a perfect heuristic, but good enough.
+      checkpoint = load_checkpoint_np(ckpt, **kw)
+      checkpoint = jax.tree.map(recover_dtype, checkpoint)
+      if "params" in checkpoint:
+        # Checkpoint with optax state (after (internal link)).
+        params = checkpoint["params"]
+      elif "opt" in checkpoint:
+        # Checkpoint with Flax optimizer.
+        params = checkpoint["opt"]["target"]
+      else:
+        # When open-sourcing, we often shared only the params directly.
+        params = checkpoint
+    else:
+      # Here we're now loading new-style tensorstore checkpoints.
+      # We can be a more efficient and load params and `key` only right away.
+      regex = f"params/{key}($|/.*)" if key else "params/.*"
+      assert "regex" not in kw, "For a custom regex, use tsload directly."
+      kw["regex"] = regex
+      checkpoint = load_checkpoint_ts(ckpt, **kw)
+      params = checkpoint["params"]
+
+  if key is not None:
+    params = tree_get(params, key)
+
+  return params
+
+
+def prefetch_scalar(it, nprefetch=1, devices=None):
+  n_loc_dev = len(devices) if devices else jax.local_device_count()
+  repl_iter = (np.ones(n_loc_dev) * i for i in it)
+  return flax_utils.prefetch_to_device(repl_iter, nprefetch, devices)
+
+
+def sigmoid_xent(*, logits, labels, reduction=True):
+  # NOTE: This implementation is stable, see these two:
+  # (internal link)
+  # https://github.com/google/jax/issues/2140
+  log_p = jax.nn.log_sigmoid(logits)
+  log_not_p = jax.nn.log_sigmoid(-logits)
+  nll = -jnp.sum(labels * log_p + (1. - labels) * log_not_p, axis=-1)
+  return jnp.mean(nll) if reduction else nll
+
+
+def bidirectional_contrastive_loss(zimg, ztxt, t, mask=None, reduction=False):
+  """Bidirectional contrastive loss (e.g. for contrastive trainer/evaluator)."""
+  # BF.FB = BB
+  logits = jnp.dot(zimg, ztxt.T) * t
+
+  if mask is not None:
+    # Set to negative infinity where mask = 0. Masked examples will disappear
+    # under softmax, and be ignored by ncorrect (NINF will never win argmax).
+    exclude = jnp.logical_not(mask)  # Now 1 if we don't want to keep.
+    exclude = jnp.logical_or(exclude[:, None], exclude[None, :])
+    logits = jnp.where(exclude, -jnp.inf, logits)
+
+  # Note: assumed t is in a good range e.g. already passed through exp/softplus.
+  l1 = -jnp.diag(jax.nn.log_softmax(logits, axis=1))  # NLL img->txt
+  l2 = -jnp.diag(jax.nn.log_softmax(logits, axis=0))  # NLL txt->img
+  l = 0.5 * (l1 + l2)
+
+  if mask is not None:
+    l = jnp.where(mask, l, 0)
+
+  redux = jnp.mean if reduction else lambda x: x
+  if reduction and mask is not None:
+    redux = lambda x: jnp.sum(x * mask) / (jnp.sum(mask) + 1e-8)
+
+  # Also return extra measurements.
+  return redux(l), {
+      "ncorrect": redux(jnp.argmax(logits, axis=1) == jnp.arange(len(logits))),
+  }
+
+
+def softmax_xent(*, logits, labels, reduction=True, kl=False, axis=-1):
+  log_p = jax.nn.log_softmax(logits, axis=axis)
+  nll = -jnp.sum(labels * log_p, axis=axis)
+  if kl:
+    nll += jnp.sum(labels * jnp.log(jnp.clip(labels, 1e-8)), axis=axis)
+  return jnp.mean(nll) if reduction else nll
+
+
+def weighted_softmax_xent(*,
+                          logits,
+                          labels,
+                          reduction=True,
+                          weights=None,
+                          label_smoothing=0.0,
+                          normalize=True):
+  """Compute weighted cross entropy.
+
+  Args:
+   logits: [batch, length, num_classes] float array.
+   labels: categorical targets [batch, length] int array.
+   reduction: reduce across batch dim.
+   weights: None or array of shape [batch, length].
+   label_smoothing: label smoothing constant, used to determine the on and off
+     values.
+   normalize: normalize each "sentence" loss by the number of tokens in it.
+
+  Returns:
+    Tuple of scalar loss and batch normalizing factor.
+  """
+  if logits.ndim != labels.ndim + 1:
+    raise ValueError("Incorrect shapes. Got shape %s logits and %s targets" %
+                     (str(logits.shape), str(labels.shape)))
+  vocab_size = logits.shape[-1]
+  confidence = 1.0 - label_smoothing
+  low_confidence = (1.0 - confidence) / (vocab_size - 1)
+  soft_targets = onehot(
+      labels, vocab_size, on_value=confidence, off_value=low_confidence)
+
+  loss = -jnp.sum(soft_targets * jax.nn.log_softmax(logits), axis=-1)
+
+  normalizing_factor = labels.shape[1]
+  if weights is not None:
+    loss = loss * weights
+    normalizing_factor = jnp.clip(weights.sum(axis=1), 2e-38)
+
+  loss = loss.sum(axis=1)
+  if normalize:
+    loss = loss / normalizing_factor
+
+  return loss.mean() if reduction else loss
+
+
+def accumulate_gradient(loss_and_grad_fn, params, images, labels, accum_steps):
+  """Accumulate gradient over multiple steps to save on memory."""
+  # See (internal link) for details and experiments.
+  if accum_steps and accum_steps > 1:
+    assert images.shape[0] % accum_steps == 0, (
+        f"Bad accum_steps {accum_steps} for batch size {images.shape[0]}")
+    step_size = images.shape[0] // accum_steps
+    l, g = loss_and_grad_fn(params, images[:step_size], labels[:step_size])
+    def acc_grad_and_loss(i, l_and_g):
+      imgs = jax.lax.dynamic_slice(images, (i*step_size, 0, 0, 0),
+                                   (step_size,) + images.shape[1:])
+      lbls = jax.lax.dynamic_slice(labels, (i*step_size, 0),
+                                   (step_size, labels.shape[1]))
+      li, gi = loss_and_grad_fn(params, imgs, lbls)
+      l, g = l_and_g
+      return (l + li, jax.tree.map(lambda x, y: x + y, g, gi))
+    l, g = jax.lax.fori_loop(1, accum_steps, acc_grad_and_loss, (l, g))
+    return jax.tree.map(lambda x: x / accum_steps, (l, g))
+  else:
+    return loss_and_grad_fn(params, images, labels)
+
+
+def itstime(step, every_n_steps, total_steps, host=None, last=True, first=True,
+            drop_close_to_last=0.25):
+  """Returns True if it's time to execute an action.
+
+  Args:
+    step: the current step representing "now".
+    every_n_steps: the action should run every this many steps.
+    total_steps: the step number of the last step of training.
+    host: host number. If provided, only run if we are this process.
+    last: whether to run on the last step or not.
+    first: whether to run on the first step or not.
+    drop_close_to_last: if a step would run, but is this close (in terms of
+      fraction of every_n_step) to the last one, skip.
+
+  Returns:
+    True if the action should be executed, False if not.
+  """
+
+  # This logic avoids running `itstime` "a few" steps before the last step.
+  # Canonical example: don't save checkpoint 2 steps before the last, and then
+  # at the last again; it's pointless and checkpoint timing will time out.
+  close_to_last = False
+  if drop_close_to_last and every_n_steps:
+    close_to_last = abs(step - total_steps) < drop_close_to_last * every_n_steps
+
+  is_host = host is None or jax.process_index() == host
+  is_step = every_n_steps and (step % every_n_steps == 0) and not close_to_last
+  is_last = every_n_steps and step == total_steps
+  is_first = every_n_steps and step == 1
+  return is_host and (is_step or (last and is_last) or (first and is_first))
+
+
+def checkpointing_timeout(writer, timeout):
+  # Make sure checkpoint writing is not a bottleneck
+  if writer is not None:
+    try:
+      # Note: `writer` is a multiprocessing.AsyncResult, and
+      # timeout is in seconds.
+      writer.get(timeout=timeout)
+    except multiprocessing.TimeoutError as e:
+      raise TimeoutError(
+          "Checkpoint writing seems to be a bottleneck. Make sure you do "
+          "not do something wrong, like writing checkpoints to a distant "
+          "cell. In a case you are OK with checkpoint writing being a "
+          "bottleneck, you can configure `ckpt_timeout` parameter") from e
+
+
+def hms(s):
+  """Format time in hours/minutes/seconds."""
+  if s < 60:
+    return f"{s:.0f}s"
+  m, s = divmod(s, 60)
+  if m < 60:
+    return f"{m:.0f}m{s:.0f}s"
+  h, m = divmod(m, 60)
+  if h < 25:
+    return f"{h:.0f}h{m:.0f}m"  # Seconds intentionally omitted.
+  d, h = divmod(h, 24)
+  return f"{d:.0f}d{h:.0f}h{m:.0f}m"  # Seconds intentionally omitted.
+
+
+class Chrono:
+  """Measures time and reports progress, hyper-specific to our train loops.
+
+  Some concepts:
+  1. This differentiates between three "types" of time:
+    - training time: the time spent on actual training (fprop/bprop/update)
+    - program time: overall time the program runs, including all overheads
+    - pause time: the chronometer can be paused (eg during evals).
+  2. This handles a "warmup": the first step is skipped for training time
+      purposes, as it includes significant compilation overheads, which distort
+      estimates.
+  3. `accum`ulates (i.e. integrates) timings, and save/load them across
+      restarts.
+  """
+
+  def __init__(self):
+    self._timing_history = collections.defaultdict(list)
+    self._measure = None
+    self._write_note = None
+
+    self.program_start_time = time.monotonic()
+    self.train_start_time = None
+    self.train_start_step = None  # When we started timing (after warmup)
+
+    self.prev_time = None
+    self.prev_step = None
+
+    self.pause_start = None
+    self.paused_time = 0
+
+    self.total_steps = None
+    self.global_bs = None
+    self.steps_per_epoch = None
+
+    self.warmup = 2  # How many calls to `tick` to skip.
+    self.load()  # Inits accum integrators.
+    self.note = "Chrono n/a"
+
+  def inform(self, *, first_step=None, total_steps=None, global_bs=None,
+             steps_per_epoch=None, measure=None, write_note=None):
+    """Provide some extra info that's only known later in the program."""
+    # The pattern of `self.x = x or self.x` allows one to call `inform` various
+    # times with various subset of information (args), as they become available.
+    # Except for `first_step` which can be 0 so is a bit more verbose.
+    self.prev_step = first_step if first_step is not None else self.prev_step
+    self.total_steps = total_steps or self.total_steps
+    self.steps_per_epoch = steps_per_epoch or self.steps_per_epoch
+    self.global_bs = global_bs or self.global_bs
+    self._measure = measure or self._measure
+    self._write_note = write_note or self._write_note
+    if self.total_steps and self.prev_step is not None:
+      self.note = (f"Steps:{self.prev_step}/{self.total_steps} "
+                   f"[{self.prev_step/self.total_steps:.1%}]")
+
+  def tick(self, step, measure=None, write_note=None):
+    """A chronometer tick."""
+    if step == self.prev_step: return  # Can happen from evals for example.
+
+    measure = measure or self._measure
+    write_note = write_note or self._write_note
+
+    now = time.monotonic()
+    measure("uptime", now - self.program_start_time)
+    self.flush_timings()
+
+    # We do always count examples, regardless of the timing-related warmup that
+    # happens a few lines below.
+    ds = step - self.prev_step  # Steps between ticks
+    self.prev_step = step
+    self.accum_examples_seen += ds * self.global_bs
+    measure("examples_seen", self.accum_examples_seen)
+    measure("progress", step / self.total_steps)
+    if self.steps_per_epoch:
+      measure("epoch", step / self.steps_per_epoch)
+
+    # We take the start as the second time `tick` is called, so we avoid
+    # measuring the overhead of compilation and don't include it in time
+    # estimates.
+    if self.warmup > 1:
+      self.warmup -= 1
+      write_note(self.note)  # This can help debugging.
+      return
+    if self.warmup == 1:
+      self.train_start_time = self.prev_time = now
+      self.train_start_step = step
+      self.accum_program_time += now - self.program_start_time
+      self.paused_time = 0  # Drop pauses that happened before timing starts.
+      self.warmup = 0
+      write_note(self.note)  # This can help debugging.
+      return
+
+    # Measurement with micro-timings of current training steps speed.
+    # Time between ticks (ignoring pause)
+    dt = now - self.prev_time - self.paused_time
+    ncores = jax.device_count()  # Global device count
+    measure("img/sec/core", self.global_bs * ds / dt / ncores)
+
+    # Accumulate (integrate) times, good for plots.
+    self.accum_train_time += dt
+    self.accum_pause_time += self.paused_time
+    self.accum_program_time += dt + self.paused_time
+
+    # Convert to, and log as, core hours.
+    core_hours = self.accum_train_time * ncores / 60 / 60
+    devtype = jax.devices()[0].device_kind
+    measure(f"core_hours_{devtype}", core_hours)
+    measure("core_hours", core_hours)  # For convenience as x-axis in sweeps.
+
+    # Progress note with "global" full-program average timings
+    # (eg in program-time minus warmup)
+    dt = now - self.train_start_time  # Time elapsed since end of warmup.
+    steps_timed = step - self.train_start_step
+    steps_todo = self.total_steps - step
+    self.note = f"Steps:{step}/{self.total_steps} [{step/self.total_steps:.1%}]"
+    self.note += f"\nWalltime:{hms(self.accum_program_time)}"
+    self.note += f" ({hms(self.accum_pause_time)} eval)"
+    self.note += f"\nETA:{hms(dt / steps_timed*steps_todo)}"
+    self.note += f"\nTotal train time:{hms(dt / steps_timed*self.total_steps)}"
+    write_note(self.note)
+
+    log_memory(measure)
+
+    self.prev_time = now
+    self.paused_time = 0
+
+  def pause(self, wait_for=()):
+    assert self.pause_start is None, "Don't pause twice."
+    jax.block_until_ready(wait_for)
+    self.pause_start = time.monotonic()
+
+  def resume(self):
+    self.paused_time += time.monotonic() - self.pause_start
+    self.pause_start = None
+
+  def save(self):
+    return dict(
+        accum_program_time=self.accum_program_time,
+        accum_train_time=self.accum_train_time,
+        accum_pause_time=self.accum_pause_time,
+        accum_examples_seen=self.accum_examples_seen,
+    )
+
+  def load(self, ckpt={}):  # pylint: disable=dangerous-default-value
+    self.accum_program_time = float(ckpt.get("accum_program_time", 0.0))
+    self.accum_train_time = float(ckpt.get("accum_train_time", 0.0))
+    self.accum_pause_time = float(ckpt.get("accum_pause_time", 0.0))
+    self.accum_examples_seen = int(ckpt.get("accum_examples_seen", 0))
+
+  @contextlib.contextmanager
+  def log_timing(self, name, *, noop=False):
+    """Use this when you time sth once per step and want instant flushing."""
+    t0 = time.monotonic()
+    yield
+    dt = time.monotonic() - t0
+    if not noop:
+      if self._measure:  # So that timed things still work in colab.
+        self._measure(name, dt)
+      logging.info("TIMING[%s]: %s", name, dt)
+      logging.flush()
+
+  @contextlib.contextmanager
+  def log_timing_avg(self, name, *, noop=False):
+    """Use this when you time sth multiple times per step (eg in a loop)."""
+    t0 = time.monotonic()
+    yield
+    dt = time.monotonic() - t0
+    if not noop:
+      self._timing_history[name].append(dt)
+      logging.info("TIMING[%s]: avg %s current %s",
+                   name, np.mean(self._timing_history[name]), dt)
+      logging.flush()
+
+  def flush_timings(self):
+    assert self._measure is not None
+    for name, times in self._timing_history.items():
+      self._measure(name, np.mean(times))
+    self._timing_history.clear()
+
+
+# Singleton to use from everywhere. https://stackoverflow.com/a/6760726/2366315
+chrono = Chrono()
+
+
+def log_memory(measure):
+  """Log a bunch of memory-related measurements."""
+  try:
+    import psutil
+  except ImportError:
+    psutil = None
+
+  if psutil is not None:
+    # Note that total != available + used, see psutil docs.
+    vmem = psutil.virtual_memory()
+    measure("y/hostmem/total", vmem.total)
+    measure("y/hostmem/available", vmem.available)
+    measure("y/hostmem/used", vmem.used)
+
+  # We show only device 0 and 1 to avoid spam. The reason to show two and not
+  # just one, if multiple are available, is because a frequent mistake is to
+  # create arrays on the default device, which is device 0.
+  for i, d in zip([0, 1], jax.local_devices()):
+    for k, v in (d.memory_stats() or {}).items():
+      measure(f"y/devmem/dev{i}/{k}", v)
+
+
+def _traverse_with_names(tree, with_inner_nodes=False):
+  """Traverses nested dicts/dataclasses and emits (leaf_name, leaf_val)."""
+  if dataclasses.is_dataclass(tree):
+    tree = flax.serialization.to_state_dict(tree)
+  # Don't output the non-leaf nodes. If the optimizer doesn't have a state
+  # the tree leaves can be Nones which was interpreted as a leaf by this
+  # function but not by the other functions (like jax.tree.map).
+  if tree is None:
+    return
+  elif isinstance(tree, Mapping):
+    keys = sorted(tree.keys())
+    for key in keys:
+      for path, v in _traverse_with_names(tree[key], with_inner_nodes):
+        yield (key + "/" + path).rstrip("/"), v
+    if with_inner_nodes:
+      yield "", tree
+  elif isinstance(tree, (list, tuple)):
+    for idx in range(len(tree)):
+      for path, v in _traverse_with_names(tree[idx], with_inner_nodes):
+        yield (str(idx) + "/" + path).rstrip("/"), v
+    if with_inner_nodes:
+      yield "", tree
+  else:
+    yield "", tree
+
+
+def tree_flatten_with_names(tree):
+  """Populates tree_flatten with leaf names.
+
+  This function populates output of tree_flatten with leaf names, using a
+  custom traversal that produces names is provided. The custom traversal does
+  NOT have to traverse tree in the same order as jax, as we take care of
+  automatically aligning jax' and custom traversals.
+
+  Args:
+    tree: python tree.
+
+  Returns:
+    A list of values with names: [(name, value), ...]
+  """
+  vals, tree_def = jax.tree.flatten(tree)
+
+  # "Fake" token tree that is use to track jax internal tree traversal and
+  # adjust our custom tree traversal to be compatible with it.
+  tokens = range(len(vals))
+  token_tree = tree_def.unflatten(tokens)
+  val_names, perm = zip(*_traverse_with_names(token_tree))
+  inv_perm = np.argsort(perm)
+
+  # Custom traverasal should visit the same number of leaves.
+  assert len(val_names) == len(vals)
+
+  return [(val_names[i], v) for i, v in zip(inv_perm, vals)], tree_def
+
+
+def tree_unflatten(names_and_vals):
+  """Reverses `tree_flatten_with_names(tree)[0]`."""
+  return recover_tree(*zip(*names_and_vals))
+
+
+def tree_map_with_names(f, tree, *rest):
+  """Like jax.tree.map but with a filter on the leaf path name.
+
+  Args:
+    f: A function with first parameter `name` (path-like "a/b/c") and remaining
+      parameters values of `tree` and `*rest` corresponding to the given `name`
+      Should return a new value for parameter `name`.
+    tree: The tree of parameters `f` should be applied to.
+    *rest: more trees of the exact same structure.
+
+  Returns:
+    A tree identical in structure to `tree` and `*rest` but with the leaves the
+    result of calling `f` on corresponding name/leaves in `tree` and `*rest`.
+  """
+  names_and_vals, tree_def = tree_flatten_with_names(tree)
+  names, vals = zip(*names_and_vals)
+  rest_vals = [list(zip(*tree_flatten_with_names(t)[0]))[1] for t in rest]
+  vals = [f(*name_and_vals) for name_and_vals in zip(names, vals, *rest_vals)]
+  return tree_def.unflatten(vals)
+
+
+def tree_map_with_regex(f, tree, regex_rules, not_f=lambda x: x, name=None):
+  """Apply jax-style tree_map based on regex rules.
+
+  Args:
+    f: a function that is being applied to every variable.
+    tree: jax tree of arrays.
+    regex_rules: a list of tuples `(pattern, args)`, where `pattern` is a regex
+      which used for variable matching and `args` are positional arguments
+      passed to `f`. If some variable is not matched, we apply `not_f` transform
+      which is id by default. If multiple patterns match, then only the first
+      rule is applied.
+    not_f: optional function which is applied to variables that do not match any
+      pattern.
+    name: a name of transform for logging purposes.
+
+  Returns:
+    a tree, transformed by `f` according to the given rules.
+  """
+  def _f(vname, v):
+    for pattern, arg in regex_rules:
+      if re.fullmatch(pattern, vname):
+        if name and jax.process_index() == 0:
+          logging.info("Applying %s to %s with %s due to `%s`",
+                       name, vname, arg, pattern)
+        return f(v, arg)
+    return not_f(v)
+  return tree_map_with_names(_f, tree)
+
+
+def tree_get(tree, name):
+  """Get an entry of pytree by flattened key name, eg a/b/c, with nice error.
+
+  Args:
+    tree: the pytree to be queried.
+    name: the path to extract from the tree, see below for examples.
+
+  Returns:
+    A few examples:
+      tree = {'a': 1, 'b': {'c': 2, 'd': 3}}
+      tree_get(tree, 'a') == 1
+      tree_get(tree, 'b/c') == 2
+      tree_get(tree, 'b') == {'c': 2, 'd': 3}
+  """
+  flattened = dict(_traverse_with_names(tree, with_inner_nodes=True))
+  try:
+    return flattened[name]
+  except KeyError as e:
+    class Msg(str):  # Reason: https://stackoverflow.com/a/70114007/2366315
+      def __repr__(self):
+        return str(self)
+    msg = "\n".join([name, "Available keys:", *flattened, ""])
+    # Turn into configdict to use its "did you mean?" error message!
+    msg = mlc.ConfigDict(flattened)._generate_did_you_mean_message(name, msg)  # pylint: disable=protected-access
+    raise KeyError(Msg(msg)) from e
+
+
+def tree_replace(tree, replacements):
+  """Renames/removes (nested) keys.
+
+  Example usage:
+
+    tree = {'a': {'b': 2, 'c': 3}, 'c': 4}
+    replacements = {
+        'a/b': 'a/b/x',  # replaces 'a/b' with 'a/b/x'
+        '.*c': 'C',      # replaces 'c' with 'C' ('a/c' is removed)
+        'C': 'D',        # replaces 'C' (which was 'c') with 'D'
+        '.*/c': None,    # removes 'a/c'
+    }
+    tree2 = rename_remove(tree, replacements)
+    assert tree2 == {'D': 4, 'a': {'b': {'x': 2}}}
+
+  Args:
+    tree: A nested dictionary.
+    replacements: Rules specifying `regex` as keys and `replacement` as values
+      to be used with `m = re.match(regex, key)` and `m.expand(replacement)`
+      for every `key` independently.
+
+      Note that:
+      1. If any rule matches with `replacement=None`, then the key is removed.
+      2. The rules are applied in order. It's possible to have multiple
+         transformations on a single key.
+
+  Returns:
+    Updated `tree` according to rules defined in `replacements`.
+  """
+  replacements = {
+      re.compile(kk): vv for kk, vv in replacements.items()
+  }
+
+  def rename(k):
+    for kk, vv in replacements.items():
+      m = kk.match(k)
+      if m:
+        k = k[:m.start()] + m.expand(vv) + k[m.end():]
+    return k
+
+  def should_remove(k):
+    return any(vv is None and kk.match(k) for kk, vv in replacements.items())
+
+  names_and_vals, _ = tree_flatten_with_names(tree)
+  names_and_vals = [
+      (rename(k), v) for k, v in names_and_vals if not should_remove(k)
+  ]
+  return tree_unflatten(names_and_vals)
+
+
+def tree_compare(tree1, tree2):
+  """Returns `(tree1_only, tree2_only, dtype_shape_mismatch)`."""
+  tree1 = flax.traverse_util.flatten_dict(tree1, sep="/")
+  tree2 = flax.traverse_util.flatten_dict(tree2, sep="/")
+  return set(tree1) - set(tree2), set(tree2) - set(tree1), {
+      k: [(v.dtype, v.shape), (tree2[k].dtype, tree2[k].shape)]
+      for k, v in tree1.items()
+      if k in tree2 and (v.dtype != tree2[k].dtype or v.shape != tree2[k].shape)
+  }
+
+
+def tree_filter(tree, mask):
+  """Returns nested dict structure with only a subset of children."""
+  # TODO: The code below only works for nested-dict and only when they
+  # have same structure. Consider relax this.
+  if not isinstance(tree, dict):
+    assert isinstance(mask, bool), f"Mask leaves must be boolean! {mask}"
+    return tree
+  assert sorted(tree.keys()) == sorted(mask.keys()), (
+      f"Keys in tree and mask are not equal! {tree.keys()} != {mask.keys()}")
+  return {k: tree_filter(v, mask[k]) for k, v in tree.items()
+          if mask[k] is not False}
+
+
+def recover_dtype(a):
+  """Numpy's `save` stores bfloat16 type as "void" type, so we recover it."""
+  if hasattr(a, "dtype") and a.dtype.type is np.void:
+    assert a.itemsize == 2, "Unknown dtype!"
+    return a.view(jax.numpy.bfloat16)
+  else:
+    return a
+
+
+def recover_tree(keys, values):
+  """Recovers a tree as a nested dict from flat names and values.
+
+  This function is useful to analyze checkpoints that are saved by our programs
+  without need to access the exact source code of the experiment. In particular,
+  it can be used to extract an reuse various subtrees of the scheckpoint, e.g.
+  subtree of parameters.
+
+  Args:
+    keys: a list of keys, where '/' is used as separator between nodes.
+    values: a list of leaf values.
+
+  Returns:
+    A nested tree-like dict.
+  """
+  tree = {}
+  sub_trees = collections.defaultdict(list)
+  for k, v in zip(keys, values):
+    if "/" not in k:
+      tree[k] = v
+    else:
+      k_left, k_right = k.split("/", 1)
+      sub_trees[k_left].append((k_right, v))
+  for k, kv_pairs in sub_trees.items():
+    k_subtree, v_subtree = zip(*kv_pairs)
+    tree[k] = recover_tree(k_subtree, v_subtree)
+  return tree
+
+
+def tssave(mngr, pytree, path, on_commit=lambda *_, **__: None):
+  """Save pytree using jax tensorstore-based checkpoint manager.
+
+  NOTE: When overwriting an existing checkpoint with a different pytree, the
+  result is, counterintuitively, the union of both, not only the new one.
+
+  Args:
+    mngr: An instance of GlobalAsyncCheckpointManager.
+    pytree: What to store; any pytree of arrays.
+    path: Where to save the pytree. Creates subfolders as needed.
+    on_commit: A callback when writing is done, see `mngr.serialize`.
+  """
+  names, vals = zip(*tree_flatten_with_names(pytree)[0])
+
+  for name in names:
+    if "~" in name:
+      raise ValueError(f"Symbol '~' is not allowed in names. Found in {name}.")
+
+  gfile.makedirs(path)
+  with jax.transfer_guard("allow"):
+    names = [name.replace("/", "~") for name in names]
+    mngr.serialize_with_paths(
+        list(vals), [os.path.join(path, name) for name in names],
+        on_commit_callback=functools.partial(on_commit, array_names=names))
+
+
+def save_checkpoint_ts(mngr, checkpoint, path, step, keep=True):
+  """Preemption-safe saving of checkpoints using tssave."""
+  # The tensorstore checkpoint format is a folder with (potentially) many files.
+  # On some file-systems, operations on these (copy, rename, delete) are slow,
+  # so we implement a flow that's both robust to pre-emptions/crashes during
+  # checkpointing and makes minimal use of these slow operations.
+
+  # The logic goes as follows. It's infaillible :)
+  #       (...if file move is atomic, which it is.)
+  # We always write the current checkpoint to a new folder, which contains the
+  # step number in its name. If we don't need to keep it indefinitely, we append
+  # "-tmp" to its name.
+  # After writing the next checkpoint, we remove the previous one if it had
+  # "-tmp" in its name.
+  # We also have a -LAST file that contains a pointer to the latest complete
+  # checkpoint. File operations are cheap to make atomic, that's why.
+
+  def _on_commit_callback(array_names):  # Runs after writing ckpt is done.
+    with gfile.GFile(f"{path}-CUR", "w") as f:
+      f.write(curr)
+
+    last = ""
+    if gfile.exists(f"{path}-LAST"):
+      with gfile.GFile(f"{path}-LAST", "r") as f:
+        last = f.read().strip()
+
+    gfile.rename(f"{path}-CUR", f"{path}-LAST", overwrite=True)
+
+    if last.endswith("-tmp"):
+      # If pre-emption happens here, some old checkpoints may not be deleted.
+      multiprocessing.pool.ThreadPool().map(
+          gfile.rmtree,
+          [f"{path}-{last}/{name}" for name in array_names])
+      gfile.rmtree(f"{path}-{last}")
+
+  # NOTE: The jax checkpoint manager automatically waits for the previous save
+  # to be finished before writing again, so we don't need to do it here.
+
+  # Always write to path with step number in it.
+  curr = f"{step:09d}{'-tmp' if not keep else ''}"
+  tssave(mngr, checkpoint, f"{path}-{curr}", _on_commit_callback)
+
+
+def load_checkpoint_ts(path, **tsload_kw):
+  """Loads a big_vision checkpoint saved by `save_checkpoint_ts`."""
+  to_load = path
+
+  try:
+    # When passing a general path (not a specific step), get the last available.
+    with gfile.GFile(f"{path}-LAST", "r") as f:
+      to_load = f"{path}-{f.read().strip()}"
+  except Exception:  # Differs based on backend, so blanket catch. pylint:disable=broad-exception-caught
+    pass
+
+  return tsload(to_load, **tsload_kw)
+
+
+def tsload(path, *, tree=None, shardings=None, regex=None):
+  """Loads tensorstore-based array-tree from disk.
+
+  If `tree` argument is provided, then array names to load and target structure
+  is derived from the tree. If `tree` is None, then array names to load are
+  derived from array filenames on the disk, and, optionally, `regex` is applied
+  to filter these names. The`tree` argument is then automatically derived from
+  array names with `recover_tree` util.
+
+  Arrays are loaded to CPU/TPU/GPU memory as specified by the `shardings`
+  argument, which is a pytree of CPU/TPU/GPU shardings (can be mixed within a
+  single pytree). `shardings` should a prefix tree of the `tree` argument. We
+  automatically broadcast `shardings` to a full `tree`. For example, a user can
+  specify `shardings=jax.sharding.SingleDeviceSharing(jax.devices('cpu')[0])`,
+  which  will be broadcasted to a full tree.
+
+  Args:
+    path: a directory where the checkpoint arrays are stored.
+    tree: a target pytree, which defines array names to load and the target tree
+      structure. If tree is None, then `tree` is inferred from the names of
+      arrays stored on the disk.
+    shardings: a prefix pytree (with respect to `tree`) of the target shardings.
+    regex: regex to filter array names from the disk, if `tree` is not provided.
+
+  Returns:
+    A pytree of loaded arrays that has the same structure as `shardings` arg.
+  """
+  if (tree is not None) and (regex is not None):
+    raise ValueError("If tree is specified, regex filtering is not allowed.")
+
+  if tree is None:
+    # Some file-systems (gs://) list folders with a trailing /, get rid of it.
+    path_names = set([p.rstrip("/").replace("~", "/")
+                      for p in gfile.listdir(path)])
+    regex = re.compile(regex) if regex is not None else re.compile(".*")
+    path_names = [p for p in path_names if regex.match(p)]
+    tree = recover_tree(path_names, [0] * len(path_names))
+
+  names_and_vals, tree_def = tree_flatten_with_names(tree)
+  names_to_load, _ = zip(*names_and_vals)
+
+  if shardings is None:
+    shardings = jax.sharding.SingleDeviceSharding(
+        jax.local_devices(backend="cpu")[0]
+    )
+  shardings = list(jax.tree.leaves(tree_broadcast(shardings, tree)))
+
+  names_to_load = [os.path.join(path, name.replace("/", "~"))
+                   for name in names_to_load]
+  specs = [array_serial.get_tensorstore_spec(n) for n in names_to_load]
+  arrays = array_serial.run_deserialization(shardings, specs, concurrent_gb=64)
+  return tree_def.unflatten(arrays)
+
+
+def steps(prefix, config, data_size=None, batch_size=None, total_steps=None,
+          default=ValueError):
+  """Gets duration named `prefix` out of `config` and converts it to steps.
+
+  Using this function to access a configuration value that denotes some kind
+  of duration (eg training time, warmup, checkpoint frequency, ...) allows the
+  duration to be specified in terms of steps, epochs, examples, or percent of
+  training time, and converts any of these into steps, such that the training
+  code only deals with steps.
+  If the result is not an integer step number, it is rounded to the nearest one.
+
+  Args:
+    prefix: The name of the duration to query. The actual config fields can
+      then be one of `prefix_steps`, `prefix_examples`, or `prefix_epochs`.
+    config: The dictionary (config) from which to read the duration.
+    data_size: The total number of training examples in one epoch.
+    batch_size: The number of examples processed per step.
+    total_steps: The total number of training steps to run.
+    default: The default value to return when no duration of the name `prefix`
+      is found in the `config`. Set to `ValueError` (the default) to raise an
+      error instead of returning a default value.
+
+  Returns:
+    The number of steps from the config, or the default value.
+
+  Raises:
+    ValueError if there is no such duration in the config and no default is set.
+  """
+  # Be helpful and make sure only match one of the following suffixes.
+  suffixes = {"steps", "examples", "epochs", "percent"}
+  matches = {
+      f"{prefix}_{s}"
+      for s in suffixes
+      if (x := config.get(f"{prefix}_{s}")) is not None and x >= 0
+  }
+  # Note that steps=0 is also a valid value (e.g. to only run evaluators).
+  assert len(matches) <= 1, f"Only one of '{matches}' should be defined."
+
+  if f"{prefix}_steps" in matches:
+    return config[f"{prefix}_steps"]
+
+  def to_integer(x):
+    # Round to nearest but always executed at least one step unless explictily
+    # asked for 0. E.g. total_epochs=0 vs total_epochs=0.0001
+    return max(1, round(x)) if x else 0
+
+  if batch_size and f"{prefix}_examples" in matches:
+    return to_integer(config[f"{prefix}_examples"] / batch_size)
+
+  if batch_size and data_size and f"{prefix}_epochs" in matches:
+    steps_per_epoch = data_size / batch_size
+    return to_integer(config[f"{prefix}_epochs"] * steps_per_epoch)
+
+  if total_steps and f"{prefix}_percent" in matches:
+    pct = config[f"{prefix}_percent"]
+    assert 0.0 <= pct <= 1.0, (  # Be helpful, since it's not obvious.
+        f"Percents should lie in [0.0, 1.0], but {prefix}_percent is {pct}")
+    return to_integer(pct * total_steps)
+
+  if default is ValueError:
+    raise ValueError(
+        f"Cannot convert {prefix} to steps, due to missing batch_size "
+        f"({batch_size}), data_size ({data_size}), total_steps ({total_steps})"
+        ", or corresponding entry in config:\n" + "\n".join(config.keys()))
+
+  return default
+
+
+def create_learning_rate_schedule(
+    total_steps, batch_size=None, data_size=None,
+    base=1.0, decay_type="stair",
+    scale_with_batchsize=False, **kw):
+  """Creates learning rate schedule, see (internal link).
+
+  Args:
+    total_steps: The total number of steps to run.
+    batch_size: The global batch-size optionally used for scaling.
+    data_size: Number of examples in the training data (for epoch conversion).
+    base: The starting learning-rate (without warmup).
+    decay_type: 'linear' or 'cosine', 'rsqrt', 'stair'.
+    scale_with_batchsize: Whether or not to scale lr automatically.
+    **kw: extra arguments specific to individual decay_types. Also contains
+      declaration of `{warmup,cooldown}_{steps,epochs,examples}` that applies
+      on top of any/all decay_type.
+
+  Returns:
+    A function learning_rate(step): float -> {"learning_rate": float}.
+  """
+
+  def to_steps(name, default=0):
+    return steps(name, kw, data_size, batch_size, total_steps, default=default)
+
+  warmup_steps = to_steps("warmup")
+  cooldown_steps = to_steps("cooldown")
+
+  # Early catch hard to backtrack errors due to warmup_steps >= total_steps,
+  # but let it run for 0 and 1 steps used to eval and debug runs.
+  assert (total_steps <= 1) or (warmup_steps < total_steps), (
+      "warmup_steps is >= total_steps")
+
+  def step_fn(step):
+    """Step to learning rate function."""
+    lr = base
+
+    # This implements the linear scaling rule following
+    # Goyal et al. at arxiv.org/abs/1706.02677.
+    # The reference batch size in literature is 256, so we scale the lr to
+    # adjust to the literature lr when bach_size changes.
+    if scale_with_batchsize:
+      lr = lr * batch_size / 256.0
+
+    progress = (step - warmup_steps) / float(total_steps - warmup_steps)
+    progress = jnp.clip(progress, 0.0, 1.0)
+    if decay_type in ("linear", "polynomial"):
+      power = kw.get("power", 1)
+      zero = kw.get("end", kw.get("linear_end", 0))
+      lr = zero + (lr - zero) * (1.0 - progress) ** power
+    elif decay_type == "cosine":
+      lr = lr * 0.5 * (1. + jnp.cos(jnp.pi * progress))
+    elif decay_type == "rsqrt":
+      # See (internal link) for details, especially how to set timescale
+      # and shift in order to continue smoothly when changing batch-size.
+      t = to_steps("timescale", default=kw.get("timescale", 10_000))
+      shift = to_steps("shift", default=kw.get("shift", 0))
+      lr = jnp.where(
+          warmup_steps <= step,
+          lr / jnp.sqrt(1 + (step + shift - warmup_steps) / t),  # In decay
+          lr / jnp.sqrt(1 + shift / t))  # In warmup.
+    elif decay_type == "stair":
+      i = jnp.searchsorted(jnp.array(kw.get("steps", [])), step + 1)
+      lr = lr * jnp.take(jnp.array([1.0] + list(kw.get("mults", []))), i)
+    else:
+      raise ValueError(f"Unknown lr type {decay_type}")
+
+    if warmup_steps:
+      lr = lr * jnp.minimum(1., step / warmup_steps)
+    if cooldown_steps:
+      lr = lr * jnp.minimum(1., (total_steps - step) / cooldown_steps)
+
+    return jnp.asarray(lr, dtype=jnp.float32)
+
+  return step_fn
+
+
+def get_mixup(rng, p):
+  """Perform mixup https://arxiv.org/abs/1710.09412."""
+  rng, rng_mixup = jax.random.split(rng)
+  a = jax.random.beta(rng_mixup, p, p)
+  a = jnp.maximum(a, 1.0 - a)  # see (internal link) for the context.
+  def _mixup(*things, **more_things):
+    mix = lambda thing: a * thing + (1 - a) * jnp.roll(thing, shift=1, axis=0)
+    return rng, *jax.tree.map(mix, (things, more_things))
+  return _mixup
+
+
+# For backwards compatability with legacy code.
+def mixup(rng, *things, p, **more_things):
+  return get_mixup(rng, p)(*things, **more_things)
+
+
+def sync():
+  """Syncs hosts and empties async computation queue."""
+  x = reshard(np.ones(jax.device_count()),
+              jax.sharding.PositionalSharding(jax.devices()))
+  jax.jit(jnp.sum)(x).block_until_ready()
+
+
+def check_and_compile_patterns(patterns):
+  """Validates and compiles a list of param-patterns.
+
+  The validation consists of checking for common mistakes, currently only that
+  the pattern does not start with a slash, because unlike FLAX, our parameter
+  names don't start with a slash.
+
+  Args:
+    patterns: a single (string) pattern (regex), or a list of patterns.
+
+  Returns:
+    A list of compiled and verified regexes.
+  """
+  if isinstance(patterns, str):
+    patterns = [patterns]
+
+  assert isinstance(patterns, (list, tuple)), patterns
+
+  def check_and_compile(pattern):
+    assert not pattern.startswith("/"), (
+        f"Big vision parameter names never start with '/': '{pattern}")
+    return re.compile(pattern)
+
+  return list(map(check_and_compile, patterns))
+
+
+def make_mask_trees(tree, patterns, *, log=None):
+  """Returns a boolean mask tree for every pattern (only first match)."""
+  compiled_patterns = check_and_compile_patterns(patterns)
+
+  def matchfirst(name, _):
+    matches = []
+    for pattern in compiled_patterns:
+      matches.append(not any(matches) and bool(pattern.fullmatch(name)))
+    if log is not None and True in matches and jax.process_index() == 0:
+      logging.info("%s: %s - matched by %s", log, name,
+                   patterns[matches.index(True)])
+    return np.array(matches)
+
+  multimask = tree_map_with_names(matchfirst, tree)
+  return [
+      jax.tree.map(lambda matches, i=idx: matches[i], multimask)
+      for idx in range(len(patterns))
+  ]
+
+
+@contextlib.contextmanager
+def profile(name, ttl=3 * 365 * 24 * 3600, noop=False):
+  if not noop:
+    sess = startstop_prof_at_steps(None, name=name, ttl=ttl)
+  yield
+  if not noop:
+    startstop_prof_at_steps(sess, name=name, ttl=ttl)
+
+
+def startstop_prof(sess, step=None, first_step=0,
+                   log_steps=1, surround=10, **kw):
+  """Runs the profiler for `surround` steps around the next `log_steps`."""
+  first_log = first_step + log_steps - (first_step % log_steps)
+  # don't start before first!
+  start = max(first_log - surround//2, first_step + 1)
+  return startstop_prof_at_steps(sess, step, start, start + surround, **kw)
+
+
+def startstop_prof_at_steps(
+    sess, step=None, first_step=None, last_step=None,
+    name="steps", ttl=3 * 365 * 24 * 3600):
+  del sess, step, first_step, last_step, name, ttl
+  pass  # TODO: implement using `jax.profiler` API. Needs workdir.
+
+
+# This is a very minimal variant for open-sourcing. Our internal code makes use
+# of multiple internal logging tools instead.
+class BigVisionMetricWriter:
+  """A class for logging metrics."""
+
+  def __init__(self, xid=-1, wid=-1, workdir=None, config=None):
+    self.step_start(0)
+    if jax.process_index() != 0: return  # Only one host shall write stuff.
+
+    self.pool = multiprocessing.pool.ThreadPool(1)  # 1 is important here.
+    self.fname = None
+    if workdir:
+      if xid != -1 and wid != -1:
+        self.fname = os.path.join(workdir,
+                                  f"big_vision_{xid}_{wid}_metrics.txt")
+      else:
+        self.fname = os.path.join(workdir, "big_vision_metrics.txt")
+      if config:
+        with gfile.GFile(os.path.join(workdir, "config.json"), "w") as f:
+          f.write(config.to_json())
+
+  def step_start(self, step):
+    self.step = step
+    self.step_metrics = {}
+
+  def measure(self, name, value):
+    """Logs the metric value."""
+    if jax.process_index() != 0: return  # Only one host shall write stuff.
+
+    # Convenience for accepting scalar np/DeviceArrays, as well as N-d single
+    # scalars, like [[[123]]] or similar, avoiding silly mistakes.
+    value = np.array(value).squeeze()
+
+    # If the value is a scalar, we keep it in mind to append a line to the logs.
+    # If it has any structure, we instead just log its shape.
+    value = float(value) if value.ndim == 0 else value.shape
+
+    logging.info(f"\u001b[35m[{self.step}]\u001b[0m {name} = {value}")
+    logging.flush()
+    self.step_metrics[name] = value
+
+    return value  # Just for convenience
+
+  def step_end(self):
+    """Ends a training step, write its full row."""
+    if not self.step_metrics: return
+
+    def write(metrics):
+      with gfile.GFile(self.fname, "a") as f:
+        f.write(json.dumps({"step": self.step, **metrics}) + "\n")
+
+    if self.fname:
+      self.pool.apply(lambda: None)  # Potentially wait for past writes.
+      self.pool.apply_async(write, (self.step_metrics,))
+
+  def close(self):
+    self.step_end()
+    if jax.process_index() == 0:
+      self.pool.close()
+      self.pool.join()
+
+
+def maybe_cleanup_workdir(workdir, cleanup, info):
+  """Potentially removes workdirs at end of run for cleanup."""
+  if not workdir:
+    return
+
+  if not cleanup:
+    info("Logs/checkpoints are in %s", workdir)
+  elif jax.process_index() == 0:
+    gfile.rmtree(workdir)
+    try:  # Only need this on the last work-unit, if already empty.
+      gfile.remove(os.path.join(workdir, ".."))
+    except tf.errors.OpError:
+      pass
+
+
+def tree_broadcast(prefix, target):
+  """Broadcasts a prefix tree to a full tree.
+
+  Input-output examples:
+  1. prefix: {"x": 10, "y": 20}
+     target: {"x": {"a": 1, "b": 2}, "y": 3}
+
+     Result: {"x": {"a": 10, "b": 10}, "y": 20}
+
+  2. prefix: 100
+     target: {"x": {"a": 1, "b": 2}, "y": 3}
+
+     Result: {"x": {"a": 100, "b": 100}, "y": 100}
+
+  3. prefix: {"x": 10}
+     target: {"x": {"a": 1, "b": 2}, "y": 3}
+
+     Result: ValueError
+
+  Args:
+    prefix: prefix pytree.
+    target: boradcast target for a prefix tree.
+
+  Returns:
+    prefix tree broadcasted to a target tree.
+  """
+  def _broadcast(leaf, subtree):
+    return jax.tree.map(lambda _: leaf, subtree)
+  return jax.tree.map(_broadcast, prefix, target)
+
+
+def reshard(tree, shardings):
+  """Take an arbitrarily* sharded pytree and shard it according to `shardings`.
+
+  This is a no-op for tree elements which are already sharded as requested.
+
+  *Arrays that are fully addressable (for example, CPU arrays) are assumed to be
+  identical (i.e. replicated) across hosts.
+
+  *It does not work if an element of `tree` is not fully-addressable, unless its
+  sharding is already consistent with the target sharding.
+  If this is needed, please ping lbeyer@ or akolesnikov@.
+
+  Args:
+    tree: a pytree of arrays.
+    shardings: a (prefix) pytree of jax array shardings.
+  Returns:
+    A pytree of global jax arrays that follows provided shardings.
+  """
+  def _make_global_arr(x, shard, shape):
+    # Avoid unnecessary copies and transfers:
+    if hasattr(x, "sharding") and x.sharding.is_equivalent_to(shard, len(shape)):  # pylint: disable=line-too-long
+      return x
+    if not getattr(x, "is_fully_addressable", True):
+      raise RuntimeError("Trying to reshard a non-fully-addressable array. "
+                         "Please see the doc-comment for detailed explanation.")
+    x = jax.device_get(x)  # Might be on local devices.
+    xs = [jax.device_put(x[s], device=d)
+          for d, s in shard.addressable_devices_indices_map(shape).items()]
+    return jax.make_array_from_single_device_arrays(shape, shard, xs)
+
+  shapes = jax.tree.map(np.shape, tree)
+  shardings = tree_broadcast(shardings, tree)
+  return jax.tree.map(_make_global_arr, tree, shardings, shapes)
+
+
+def put_cpu(x):
+  """Places array/pytree on a CPU device."""
+  return jax.device_put(x, jax.local_devices(backend="cpu")[0])
+
+
+def make_fsarray_from_local_slice(local_slice, global_devices):
+  """Create a fully-sharded global device array from local host arrays.
+
+  Args:
+    local_slice: Something convertible to a numpy array (eg also TF tensors)
+      that is this host's slice of the global array.
+    global_devices: The list of global devices. Needed for consistent ordering.
+
+  Returns:
+    The global on-device array which consists of all local slices stacked
+    together in the order consistent with the devices.
+  """
+  mesh = jax.sharding.Mesh(global_devices, ("devices",))
+  sharding = jax.sharding.NamedSharding(
+      mesh, jax.sharding.PartitionSpec("devices"))
+  local_ds = mesh.local_devices
+
+  x = np.asarray(memoryview(local_slice))  # No-copy: http://(internal link)
+  xs = jax.device_put(np.split(x, len(local_ds), axis=0), local_ds)
+
+  global_shape = (x.shape[0] * jax.process_count(), *x.shape[1:])
+  return jax.make_array_from_single_device_arrays(global_shape, sharding, xs)
+
+
+def get_local_slice_from_fsarray(global_array):
+  """Return numpy array for the host-local slice of fully-sharded array.
+
+  Args:
+    global_array: JAX array, globally sharded on devices across hosts.
+
+  Returns:
+    NumPy array that holds the part of `global_array` that is held by the
+    devices on the host that calls this function.
+  """
+  # For now, for simplicity, we only implement slicing along the first axis.
+  for shard in global_array.addressable_shards:
+    assert all(idx == slice(None) for idx in shard.index[1:]), (
+        f"global_array is sharded along non-first dimensions:\n{shard.index}")
+
+  # Get the shards back in the same order in which the global array was created
+  # in the first place. This makes sure it's consistent with other things in the
+  # batch, for example (assuming the whole batch is consistent).
+  m = {s.device: s for s in global_array.addressable_shards}
+  local_shards = [m[d] for d in global_array.sharding.mesh.local_devices]
+  return np.concatenate([jax.device_get(s.data) for s in local_shards], axis=0)
+
+
+def assert_local_slices_same(*global_arrays):
+  """Check whether all `global_arrays` have local slices at the same indices."""
+  slices = [
+      tuple(
+          tuple((idx.start, idx.end, idx.step) for idx in s.index)
+          for s in a.addressable_shards)
+      for a in global_arrays]
+  assert len(set(slices)) == 1, f"Not all slices are the same: {slices}"
+
+
+# TODO: remove this logic when the
+# issue is github fixed https://github.com/google/jax/issues/15600.
+def jit_cpu(**extra_kwargs):
+  def _decorator(fun):
+    def _wrapped(*args, **kwargs):
+      sh = jax.sharding.SingleDeviceSharding(
+          jax.local_devices(backend="cpu")[0]
+      )
+      return jax.jit(fun, **extra_kwargs, out_shardings=sh)(*args, **kwargs)
+    return _wrapped
+  return _decorator
+
+
+def create_device_mesh(
+    config_mesh,
+    *,
+    allow_split_physical_axes=False,
+):
+  """Returns a JAX device mesh.
+
+  Args:
+    config_mesh: A list of tuples of (axis_name, axis_size). It is advised to
+      sort the axis in increasing order of network communication intensity.
+    allow_split_physical_axes: Whether to allow splitting physical axes.
+  """
+  devices = jax.devices()
+  mesh_axes, mesh_size = tuple(zip(*config_mesh))
+  # Because jax.utils do not support `-1` shape size.
+  mesh_size = np.array(devices).reshape(mesh_size).shape
+  device_mesh = mesh_utils.create_device_mesh(
+      mesh_size,
+      devices=devices,
+      allow_split_physical_axes=allow_split_physical_axes)
+  return jax.sharding.Mesh(device_mesh, mesh_axes)

Tipsomaly/model/big_vision/utils_test.py

@@ -0,0 +1,360 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for utils."""
+
+from functools import partial
+import os
+
+from absl.testing import parameterized
+from big_vision import utils
+import chex
+import flax
+import jax
+from jax.experimental.array_serialization import serialization as array_serial
+import jax.numpy as jnp
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.io import gfile
+
+
+NDEV = 4
+
+
+def setUpModule():
+  chex.set_n_cpu_devices(NDEV)
+
+
+class PadShardUnpadTest(chex.TestCase, tf.test.TestCase):
+  BATCH_SIZES = [NDEV, NDEV + 1, NDEV - 1, 5 * NDEV, 5 * NDEV + 1, 5 * NDEV - 1]
+  DTYPES = [np.float32, np.uint8, jax.numpy.bfloat16, np.int32]
+
+  def tearDown(self):
+    chex.clear_trace_counter()
+    super().tearDown()
+
+  @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)
+  def test_basics(self, dtype, bs):
+    # Just tests that basic calling works without exploring caveats.
+    @partial(utils.pad_shard_unpad, static_argnums=())
+    def add(a, b):
+      return a + b
+
+    x = jnp.arange(bs, dtype=dtype)
+    y = add(x, 10 * x)
+    chex.assert_type(y.dtype, x.dtype)
+    np.testing.assert_allclose(np.float64(y), np.float64(x + 10*x))
+
+  @parameterized.parameters(DTYPES)
+  def test_min_device_batch_avoids_recompile(self, dtype):
+    @partial(utils.pad_shard_unpad, static_argnums=())
+    @jax.jit
+    @chex.assert_max_traces(n=1)
+    def add(a, b):
+      return a + b
+
+    chex.clear_trace_counter()
+
+    for bs in self.BATCH_SIZES:
+      x = jnp.arange(bs, dtype=dtype)
+      y = add(x, 10 * x, min_device_batch=9)  # pylint: disable=unexpected-keyword-arg
+      chex.assert_type(y.dtype, x.dtype)
+      np.testing.assert_allclose(np.float64(y), np.float64(x + 10*x))
+
+  @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)
+  def test_static_argnum(self, dtype, bs):
+    @partial(utils.pad_shard_unpad, static_argnums=(1,))
+    def add(a, b):
+      return a + b
+
+    x = jnp.arange(bs, dtype=dtype)
+    y = add(x, dtype(10))
+    chex.assert_type(y.dtype, x.dtype)
+    np.testing.assert_allclose(np.float64(y), np.float64(x + 10))
+
+  @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)
+  def test_static_argnames(self, dtype, bs):
+    # In this test, leave static_argnums at the default value too, in order to
+    # test the default/most canonical path where `params` are the first arg.
+    @partial(utils.pad_shard_unpad, static_argnames=('b',))
+    def add(params, a, *, b):
+      return params * a + b
+
+    x = jnp.arange(bs, dtype=dtype)
+    y = add(dtype(5), x, b=dtype(10))
+    chex.assert_type(y.dtype, x.dtype)
+    np.testing.assert_allclose(np.float64(y), np.float64(5 * x + 10))
+
+
+class TreeTest(tf.test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.d1 = {'w1': 1, 'w2': 2, 'w34': (3, 4)}
+    self.d1_flat = [1, 2]
+    self.d1_flat_jax = jax.tree.flatten(self.d1)[0]
+    self.d1_named_flat = [('w1', 1), ('w2', 2), ('w34/0', 3), ('w34/1', 4)]
+    self.d1_named_flat_jax = [('w1', 1), ('w2', 2), ('w34/0', 3), ('w34/1', 4)]
+
+    self.d2 = {'conv1': {'kernel': 0, 'bias': 1},
+               'conv2': {'kernel': 2, 'bias': 3}}
+    self.d2_flat = [1, 0, 3, 2]
+    self.d2_flat_jax = jax.tree.flatten(self.d2)[0]
+    self.d2_named_flat = [('conv1/bias', 1), ('conv1/kernel', 0),
+                          ('conv2/bias', 3), ('conv2/kernel', 2)]
+    self.d2_named_flat_jax = [('conv1/bias', 1), ('conv1/kernel', 0),
+                              ('conv2/bias', 3), ('conv2/kernel', 2)]
+    self.d2_named_flat_inner = [
+        ('conv1/bias', 1), ('conv1/kernel', 0), ('conv1', self.d2['conv1']),
+        ('conv2/bias', 3), ('conv2/kernel', 2), ('conv2', self.d2['conv2']),
+        ('', self.d2),
+    ]
+
+    # This is a very important testcase that checks whether we correctly
+    # recover jax' traversal order, even though our custom traversal may not
+    # be consistent with jax' traversal order. In particular, jax traverses
+    # FlaxStruct in the order of attribute definition, while our custom
+    # traversal is alphabetical.
+    @flax.struct.dataclass
+    class FlaxStruct():
+      v3: float
+      v2: int
+      v1: str
+    self.d3 = {'a': 0, 'flax': FlaxStruct(2.0, 1, 's')}
+    self.d3_flat = [0, 1, 2.0, 's']
+    self.d3_flat_jax = jax.tree.flatten(self.d3)[0]
+    self.d3_named_flat = [
+        ('a', 0), ('flax/v1', 's'), ('flax/v2', 1), ('flax/v3', 2.0)]
+    self.d3_named_flat_jax = [
+        ('a', 0), ('flax/v3', 2.0), ('flax/v2', 1), ('flax/v1', 's')]
+
+  def test_traverse_with_names(self):
+    names_and_vals = list(utils._traverse_with_names(self.d1))
+    self.assertEqual(names_and_vals, self.d1_named_flat)
+
+    names_and_vals = list(utils._traverse_with_names(self.d2))
+    self.assertEqual(names_and_vals, self.d2_named_flat)
+
+    names_and_vals = list(utils._traverse_with_names(
+        self.d2, with_inner_nodes=True))
+    self.assertEqual(names_and_vals, self.d2_named_flat_inner)
+
+    names_and_vals = list(utils._traverse_with_names(self.d3))
+    self.assertEqual(names_and_vals, self.d3_named_flat)
+
+  def test_tree_flatten_with_names(self):
+    names_and_vals = utils.tree_flatten_with_names(self.d1)[0]
+    self.assertEqual(names_and_vals, self.d1_named_flat_jax)
+    self.assertEqual([x for _, x in names_and_vals], self.d1_flat_jax)
+
+    names_and_vals = utils.tree_flatten_with_names(self.d2)[0]
+    self.assertEqual(names_and_vals, self.d2_named_flat_jax)
+    self.assertEqual([x for _, x in names_and_vals], self.d2_flat_jax)
+
+    names_and_vals = utils.tree_flatten_with_names(self.d3)[0]
+    self.assertEqual(names_and_vals, self.d3_named_flat_jax)
+    self.assertEqual([x for _, x in names_and_vals], self.d3_flat_jax)
+
+  def test_tree_map_with_names(self):
+    d1 = utils.tree_map_with_names(
+        lambda name, x: -x if 'w2' in name else x, self.d1)
+    self.assertEqual(d1, {'w1': 1, 'w2': -2, 'w34': (3, 4)})
+
+    d1 = utils.tree_map_with_names(
+        lambda name, x1, x2: x1 + x2 if 'w2' in name else x1, self.d1, self.d1)
+    self.assertEqual(d1, {'w1': 1, 'w2': 4, 'w34': (3, 4)})
+
+  def test_recover_tree(self):
+    keys = ['a/b', 'a/c/x', 'a/c/y', 'd']
+    values = [0, 1, 2, 3]
+    self.assertEqual(utils.recover_tree(keys, values),
+                     {'a': {'b': 0, 'c': {'x': 1, 'y': 2}}, 'd': 3})
+
+  def test_make_mask_trees(self):
+    F, T = False, True  # pylint: disable=invalid-name
+    tree = {'a': {'b': 0, 'x': 1}, 'b': {'x': 2, 'y': 3}}
+    msk1 = {'a': {'b': F, 'x': T}, 'b': {'x': T, 'y': F}}
+    msk2 = {'a': {'b': F, 'x': F}, 'b': {'x': F, 'y': T}}
+    # Note that 'b' matches '^b' only and not '.*/b'.
+    # Also note that "b/x" is matched by rule 1 only (because it comes first).
+    self.assertEqual(
+        utils.make_mask_trees(tree, ('.*/x', 'b/.*')), [msk1, msk2])
+
+  def test_tree_get(self):
+    tree = {'a': {'b': 0, 'x': 1}, 'b': {'x': 2, 'y': 3}}
+    self.assertEqual(utils.tree_get(tree, 'a/b'), 0)
+    self.assertEqual(utils.tree_get(tree, 'a/x'), 1)
+    self.assertEqual(utils.tree_get(tree, 'b/x'), 2)
+    self.assertEqual(utils.tree_get(tree, 'b/y'), 3)
+    self.assertEqual(utils.tree_get(tree, 'a'), tree['a'])
+    self.assertEqual(utils.tree_get(tree, 'b'), tree['b'])
+
+  def test_tree_replace(self):
+    tree = {'a': {'b': 2, 'c': 3}, 'c': 4}
+    replacements = {
+        'a/b': 'a/b/x',  # replaces 'a/b' with 'a/b/x'
+        '.*c': 'C',      # replaces 'c' with 'C' ('a/c' is removed)
+        'C': 'D',        # replaces 'C' (which was 'c') with 'D'
+        '.*/c': None,    # removes 'a/c'
+    }
+    tree2 = utils.tree_replace(tree, replacements)
+    self.assertEqual(tree2, {'D': 4, 'a': {'b': {'x': 2}}})
+
+  def test_tree_compare(self):
+    tree1_only, tree2_only, dtype_shape_mismatch = utils.tree_compare(
+        {'a': {'b': jnp.array(2), 'c': jnp.array(3)}},
+        {'a': {'B': jnp.array(2), 'c': jnp.array(3.)}},
+    )
+    self.assertEqual(tree1_only, {'a/b'})
+    self.assertEqual(tree2_only, {'a/B'})
+    self.assertEqual(
+        dtype_shape_mismatch,
+        {'a/c': [(jnp.dtype('int32'), ()), (jnp.dtype('float32'), ())]})
+
+
+class StepConversionTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ('nice_steps', 1000, None, None, dict(foo_steps=3), 3),
+      ('nice_epochs', 1000, 100, None, dict(foo_epochs=3), 30),
+      ('nice_examples', None, 100, None, dict(foo_examples=300), 3),
+      ('nice_percent', None, None, 10, dict(foo_percent=0.30), 3),
+      ('ignore_neg', 1000, 100, 10, dict(foo_steps=-1, foo_epochs=-1,
+                                         foo_examples=-1, foo_percent=0.30), 3),
+      ('zero_steps', None, None, 10, dict(foo_percent=0.0), 0),
+      ('offbyone_steps', 1001, None, None, dict(foo_steps=3), 3),
+      ('offbyone_epochs', 1001, 100, None, dict(foo_epochs=3), 30),
+      ('offbyone_examples', None, 101, None, dict(foo_examples=300), 3),
+      ('offbyone_percent', None, None, 11, dict(foo_percent=0.30), 3),
+  )
+  def test_steps(self, data_size, batch_size, total, cfg, expected):
+    # Correct default usage:
+    step = utils.steps('foo', cfg, data_size=data_size, batch_size=batch_size,
+                       total_steps=total)
+    self.assertEqual(step, expected)
+
+    # Inexitent entry:
+    with self.assertRaises(ValueError):
+      step = utils.steps('bar', cfg, data_size=data_size, batch_size=batch_size,
+                         total_steps=total)
+    step = utils.steps('bar', cfg, data_size=data_size, batch_size=batch_size,
+                       total_steps=total, default=1234)
+    self.assertEqual(step, 1234)
+
+
+class CreateLearningRateScheduleTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.named_parameters(
+      ('linear', 'linear', {}, 13, .5),
+      ('polynomial', 'polynomial', {'end': .1, 'power': 2}, 13, .325),
+      ('cosine', 'cosine', {}, 13, .5),
+      ('rsqrt', 'rsqrt', {'timescale': 1}, 13, 0.3333333),
+      ('stair_5', 'stair', {'steps': [10], 'mults': [.5]}, 5, 1.),
+      ('stair_10', 'stair', {'steps': [10], 'mults': [.5]}, 10, .5),
+      ('warmup_before', 'rsqrt', {'timescale': 1}, 3, .6),
+      ('cooldown_after', 'rsqrt', {'timescale': 1}, 20, .05),
+  )
+  def test_schedule(self, decay_type, extra_kwargs, step, expected_lr):
+    lr_fn = utils.create_learning_rate_schedule(
+        total_steps=21,
+        batch_size=512,
+        base=.5,
+        decay_type=decay_type,
+        scale_with_batchsize=True,
+        warmup_steps=5,
+        cooldown_steps=5,
+        **extra_kwargs)
+    lr = lr_fn(step)
+    self.assertAlmostEqual(lr, expected_lr)
+
+
+class CheckpointTest(tf.test.TestCase):
+
+  def setup(self):
+    gacm = array_serial.GlobalAsyncCheckpointManager()
+
+    save_path = os.path.join(self.create_tempdir('workdir'), 'checkpoint.bv')
+    x = utils.put_cpu(np.array([1, 2, 3, 4]))
+    y = utils.put_cpu(np.array([5, 6, 7, 8]))
+    ckpt = {'x': x, 'y': {'z': y}}
+
+    sharding = jax.sharding.SingleDeviceSharding(
+        jax.local_devices(backend='cpu')[0]
+    )
+    shardings = jax.tree.map(lambda _: sharding, ckpt)
+
+    return gacm, save_path, ckpt, shardings
+
+  def test_save_and_load(self):
+    gacm, save_path, ckpt, shardings = self.setup()
+    step = 100
+    utils.save_checkpoint_ts(gacm, ckpt, save_path, step, keep=True)
+    gacm.wait_until_finished()
+    ckpt_loaded = utils.load_checkpoint_ts(save_path,
+                                           tree=ckpt, shardings=shardings)
+    chex.assert_trees_all_equal(ckpt_loaded, ckpt)
+
+    save_path_step = f'{save_path}-{step:09d}'
+    ckpt_loaded_step = utils.tsload(save_path_step, shardings=shardings)
+    chex.assert_trees_all_equal(ckpt_loaded_step, ckpt)
+
+  def test_save_and_partial_load(self):
+    gacm, save_path, ckpt, shardings = self.setup()
+    utils.save_checkpoint_ts(gacm, ckpt, save_path, step=100)
+    gacm.wait_until_finished()
+    _ = shardings.pop('x'), ckpt.pop('x')
+    ckpt_loaded = utils.load_checkpoint_ts(save_path,
+                                           tree=ckpt, shardings=shardings)
+    chex.assert_trees_all_equal(ckpt_loaded, ckpt)
+
+  def test_save_and_cpu_load(self):
+    gacm, save_path, ckpt, _ = self.setup()
+    utils.save_checkpoint_ts(gacm, ckpt, save_path, step=100)
+    gacm.wait_until_finished()
+    ckpt_loaded = utils.load_checkpoint_ts(save_path)
+    chex.assert_trees_all_equal(ckpt_loaded, ckpt)
+
+  def test_save_and_partial_cpu_load(self):
+    gacm, save_path, ckpt, _ = self.setup()
+    utils.save_checkpoint_ts(gacm, ckpt, save_path, step=100)
+    gacm.wait_until_finished()
+    ckpt.pop('y')
+    ckpt_loaded = utils.load_checkpoint_ts(save_path, regex='x.*')
+    chex.assert_trees_all_equal(ckpt_loaded, ckpt)
+
+  def test_keep_deletes(self):
+    def x(tree, factor):  # x as in "times" for multiplying.
+      return jax.tree.map(lambda a: a * factor, tree)
+
+    gacm, save_path, ckpt, _ = self.setup()
+    utils.save_checkpoint_ts(gacm, ckpt, save_path, step=100, keep=False)
+    utils.save_checkpoint_ts(gacm, x(ckpt, 2), save_path, step=200, keep=True)
+    utils.save_checkpoint_ts(gacm, x(ckpt, 3), save_path, step=300, keep=False)
+    gacm.wait_until_finished()
+    ckpt_loaded_200 = utils.tsload(f'{save_path}-{200:09d}')
+    chex.assert_trees_all_equal(ckpt_loaded_200, x(ckpt, 2))
+    ckpt_loaded_300 = utils.tsload(f'{save_path}-{300:09d}-tmp')
+    chex.assert_trees_all_equal(ckpt_loaded_300, x(ckpt, 3))
+    ckpt_loaded_last = utils.load_checkpoint_ts(save_path)
+    chex.assert_trees_all_equal(ckpt_loaded_last, x(ckpt, 3))
+    with self.assertRaises(Exception):  # Can different types depending on fs.
+      _ = utils.tsload(f'{save_path}-{100:09d}')
+    # Test that ckpt@100 was deleted
+    self.assertFalse(gfile.exists(f'{save_path}-{100:09d}-tmp'))
+
+
+if __name__ == '__main__':
+  tf.test.main()

Tipsomaly/model/omaly/__init__.py

@@ -0,0 +1,2 @@
+from .text_encoder import text_encoder
+from .vision_encoder import vision_encoder

Tipsomaly/model/omaly/fixed_prompts.py

@@ -0,0 +1,54 @@
+def generate_prompt_templates(prompt_type):
+    if prompt_type == 'medical':
+        prompt_normal = [
+            'normal {}',
+            'intact {}',
+            '{} with uniform structure',
+            '{} showing clear tissue',
+            '{} with normal anatomy',
+            '{} showing no distortion',
+            '{} with symmetric appearance',
+            '{} looking normal',
+            '{} with even texture',
+            '{} with regular shape',
+        ]
+        prompt_abnormal = [
+            'abnormal {}',
+            '{} with spot',
+            '{} with abnormality',
+            'diseased {}',
+            '{} showing distortion',
+            '{} with irregular area',
+            '{} with irregular shape',
+            '{} with uneven texture',
+        ]
+        # Prompt templates adapted for medical images
+        prompt_templates = [
+            'a medical image of a {}.',
+            'a medical image of the {}.',
+            'a diagnostic scan of a {}.',
+            'a diagnostic scan of the {}.',
+            'a slice showing a {}.',
+            'a slice showing the {}.',
+            'a scan of the {}.',
+            'a clinical brain scan of a {}.'
+        ]      
+
+    elif prompt_type == 'object_agnostic':
+        prompt_normal = ['{}']
+        # prompt_normal = ['normal {}']
+        prompt_abnormal = ['damaged {}']
+        prompt_templates = ['{}']
+        # prompt_templates = ['a photo of a {}']
+        
+    elif prompt_type == 'industrial':
+        prompt_normal = ['{}', 'flawless {}', 'perfect {}', 'unblemished {}', '{} without flaw',
+                            '{} without defect',
+                            '{} without damage']
+        prompt_abnormal = ['damaged {}', 'broken {}', '{} with flaw', '{} with defect', '{} with damage']
+        prompt_templates = ['a bad photo of a {}.',
+                            'a low resolution photo of the {}.',
+                            'a bad photo of the {}.',
+                            'a cropped photo of the {}.',
+                            ]        
+    return prompt_normal, prompt_abnormal, prompt_templates

Tipsomaly/model/omaly/text_encoder.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+# NOTE: Since tips is inside model/, your import in model/omaly/text_encoder.py should be:
+from model.tips.text_encoder import TextEncoder as BaseTextEncoder
+from .fixed_prompts import generate_prompt_templates
+import jax.numpy as jnp
+import numpy as np
+from tqdm import tqdm
+import math
+
+def jax_to_torch(x):
+    return torch.from_numpy(np.array(x))
+
+class text_encoder(nn.Module):
+    def __init__(self, tokenizer, bb_text_encoder, bb_type, text_embd_dim, MAX_LEN, prompt_learn_method='none', prompt_type='industrial', n_prompt=8, n_deep=0, d_deep=0):
+        super(text_encoder, self).__init__()  
+        self.tokenizer = tokenizer
+        self._encoder = bb_text_encoder
+        self.model = bb_type
+        self.MAX_LEN = MAX_LEN
+        self.prompt_learn_method = prompt_learn_method
+        self.n_deep_tokens = n_deep
+        self.d_deep_tokens = d_deep
+        self.n_prompt = n_prompt
+        self.prompt_type = prompt_type
+        self.text_embd_dim = text_embd_dim
+        
+        self.prompt_normal, self.prompt_abnormal, self.prompt_templates = generate_prompt_templates(self.prompt_type)
+
+        self.prompt_state = [self.prompt_normal, self.prompt_abnormal]
+    
+        if self.n_deep_tokens > 0 and self.d_deep_tokens > 0:
+            self.deep_parameters = torch.nn.ParameterList([torch.nn.Parameter(\
+                                    torch.randn(self.n_deep_tokens, text_embd_dim) * 0.02) \
+                                    for _ in range(self.d_deep_tokens)])
+        else:
+            self.deep_parameters = None
+    
+        if not self.prompt_learn_method == 'none':
+            self.normal_prompt = torch.nn.Parameter(torch.randn(self.n_prompt, text_embd_dim) * 0.02)  # Learnable prompt for normal text description with std of 0.02
+            self.abnormal_prompt = torch.nn.Parameter(torch.randn(self.n_prompt, text_embd_dim) * 0.02)  # Learnable prompt for abnormal text description with std of 0.02
+            self.learnable_prompts = torch.nn.ParameterList([self.normal_prompt, self.abnormal_prompt])
+        
+        
+    def forward(self, texts, device, learned=False):
+        text_feature_list = []
+
+        for indx, text in enumerate(texts):
+            text_features = self.encode_text(text, device, learned)
+            text_feature_list.append(text_features)
+
+        text_features = torch.stack(text_feature_list, dim=0)
+        text_features = F.normalize(text_features, dim=2)
+        return text_features
+
+    def encode_text(self, text, device, learned=False):
+        text_features = []
+        for i in range(len(self.prompt_state)):
+            learnables = self.learnable_prompts[i] if learned and not self.prompt_learn_method == 'none' else None
+            deep_parameters = self.deep_parameters if learned else None
+            
+            text = text.replace('-', ' ')
+            prompted_state = [state.format(text) for state in self.prompt_state[i]]
+            prompted_sentence = []
+            for s in prompted_state:
+                for template in self.prompt_templates:
+                    prompted_sentence.append(template.format(s))
+            
+            if self.model == 'tips':
+                # NOTE: replace the class based prompt learning concatenated to the templates with only 2 sentences
+                text_ids, text_paddings = self.tokenizer.tokenize(prompted_sentence, max_len=self.MAX_LEN)
+                class_embeddings = self._encoder(text_ids.to(device), text_paddings.to(device), learnables, self.prompt_learn_method, deep_parameters, device)
+
+                # NOTE: Avoid in-place /=, +=, _add() on tensor which are actively gradiented 
+                class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True).clamp(min=1e-3)
+                class_embedding = class_embeddings.mean(dim=0)
+                class_embedding = class_embedding / class_embedding.norm(dim=-1, keepdim=True)
+            
+            elif self.model == 'siglip2':
+                num_prompts = len(prompted_sentence)
+                if num_prompts == 0:
+                    class_embedding = torch.zeros(getattr(self, "expected_text_dim", 512), device=device)
+                    text_features.append(class_embedding)
+                    continue
+
+                batch_size = 20
+                batch_ranges = range(0, num_prompts, batch_size)
+                total_batches = math.ceil(num_prompts / batch_size)
+
+                # accumulator for sum in JAX (None -> initialize on first batch)
+                sum_jax = None
+
+                # optional tqdm over batches
+                batch_iter = tqdm(batch_ranges, total=total_batches,
+                                desc=f"Encoding text {i+1}/{len(self.prompt_state)} ({self.model})",
+                                leave=False)
+
+                for start in batch_iter:
+                    batch_sentences = prompted_sentence[start:start + batch_size]
+                    # tokenizer expects a list of strings (batch)
+                    txts = self.tokenizer(batch_sentences)
+
+                    # encoder returns jax arrays; ztxt shape -> (batch, dim)
+                    _, ztxt, out = self._encoder(txts)
+
+                    # sum along batch axis in JAX to accumulate raw (unnormalized) vectors
+                    batch_sum = jnp.sum(ztxt, axis=0)  # shape (dim,)
+
+                    if sum_jax is None:
+                        sum_jax = batch_sum
+                    else:
+                        sum_jax = sum_jax + batch_sum
+
+                # now compute mean in JAX
+                mean_jax = sum_jax / float(num_prompts)
+                mean_jax = mean_jax / (jnp.linalg.norm(mean_jax, axis=-1, keepdims=True) + 1e-8)
+
+                # convert final normalized mean to torch once
+                class_embedding = jax_to_torch(mean_jax)  # shape: (dim,)
+            elif self.model == 'siglip2-hf':
+                ids = self.tokenizer(text=prompted_sentence, padding="max_length", max_length=self.MAX_LEN, return_tensors="pt")
+                ztxt = self._encoder(ids['input_ids'].to(device), learnable_prompts=learnables, learning_method=self.prompt_learn_method).pooler_output
+                ztxt = ztxt / ztxt.norm(dim=-1, keepdim=True).clamp(min=1e-3)
+                ztxt = ztxt.mean(dim=0)
+                class_embedding = ztxt / ztxt.norm(dim=-1, keepdim=True)
+            # elif self.model == 'siglip2-hf': # NOTE: Use this block if you lack enough memory for encoding text input
+            #     num_prompts = len(prompted_sentence)
+            #     if num_prompts == 0:
+            #         class_embedding = torch.zeros(getattr(self, "expected_text_dim", 512), device=device)
+            #         text_features.append(class_embedding)
+            #         continue
+
+            #     batch_size = 20
+            #     batch_ranges = range(0, num_prompts, batch_size)
+            #     total_batches = math.ceil(num_prompts / batch_size)
+
+            #     sum_torch = None
+
+            #     # optional tqdm over batches
+            #     batch_iter = tqdm(batch_ranges, total=total_batches,
+            #                     desc=f"Encoding text {i+1}/{len(self.prompt_state)} ({self.model})",
+            #                     leave=False)
+                
+            #     for start in batch_iter:
+            #         batch_sentences = prompted_sentence[start:start + batch_size]
+            #         ids = self.tokenizer(text=batch_sentences, padding="max_length", max_length=self.MAX_LEN, return_tensors="pt")
+            #         ztxt = self._encoder(ids['input_ids'].to(device), learnable_prompts=learnables, learning_method=self.prompt_learn_method).pooler_output
+            #         batch_sum = torch.sum(ztxt, dim=0)  # shape (dim,)
+
+            #         if sum_torch is None:
+            #             sum_torch = batch_sum
+            #         else:
+            #             sum_torch = sum_torch + batch_sum
+                
+            #     mean_torch = sum_torch / float(num_prompts)
+            #     mean_torch = mean_torch / (torch.norm(mean_torch, dim=-1, keepdim=True) + 1e-8)
+            #     class_embedding = mean_torch
+
+            text_features.append(class_embedding)
+
+        text_features = torch.stack(text_features, dim=0)
+        return text_features

Tipsomaly/model/omaly/vision_encoder.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import jax.numpy as jnp
+import numpy as np
+
+def jax_to_torch(x):
+    return torch.from_numpy(np.array(x))
+    
+# The backbone can be SigLIP2 or TIPS
+class vision_encoder(nn.Module):
+    def __init__(self, bb_vision_encoder, bb_type):
+        super(vision_encoder, self).__init__()
+        self._encoder = bb_vision_encoder
+        self.model = bb_type
+
+    def _forward_tips(_encoder, images):
+        outputs = _encoder(images)
+        
+        first_cls_token = outputs[0]
+        first_cls_token = first_cls_token / first_cls_token.norm(dim=-1, keepdim=True).clamp(min=1e-3)
+        
+        second_cls_token = outputs[1]
+        second_cls_token = second_cls_token / second_cls_token.norm(dim=-1, keepdim=True).clamp(min=1e-3)
+
+        spatial_tokens = outputs[2]
+        spatial_tokens = spatial_tokens / spatial_tokens.norm(dim=-1, keepdim=True).clamp(min=1e-3)
+        
+        return first_cls_token, second_cls_token, spatial_tokens
+    
+    def _forward_siglip2(_encoder, images):
+        _, _, out = _encoder(images)
+        return jax_to_torch(out['img/normalized']), jax_to_torch(out['img/normalized']), jax_to_torch(out['img/2d_normalized'])
+    
+    def _forward_siglip2_hf(_encoder, images):
+        image_features = _encoder(images)
+        return image_features.pooler_output, image_features.pooler_output, image_features.last_hidden_state
+
+    def forward(self, images):
+        if self.model == 'siglip2':
+            return vision_encoder._forward_siglip2(self._encoder, images)
+        elif self.model == 'tips':
+            return vision_encoder._forward_tips(self._encoder, images)
+        elif self.model == "siglip2-hf":
+            return vision_encoder._forward_siglip2_hf(self._encoder, images)
+         
+    # outputs /= outputs.norm(dim=-1, keepdim=True).clamp(min=1e-3)
+    # first_cls_token, second_cls_token, spatial_tokens = outputs[:, 0], outputs[:, 1], outputs[:, 2:]

Tipsomaly/model/siglip2/siglip2_prompt_learnable.py

@@ -0,0 +1,183 @@
+from typing import Optional
+import torch
+import torch.nn as nn
+from transformers.models.siglip.modeling_siglip import SiglipTextConfig, SiglipTextEmbeddings, SiglipTextTransformer, SiglipTextModel, SiglipPreTrainedModel
+from transformers.utils import (
+    auto_docstring,
+    TransformersKwargs,
+    can_return_tuple,
+)
+
+from transformers.utils.generic import check_model_inputs
+from transformers.processing_utils import Unpack
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPooling,
+)
+
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
+
+class SiglipTextEmbeddingsWithPromptLearning(SiglipTextEmbeddings):
+    """
+    Extends SiglipTextEmbeddings with learnable prompt tokens.
+    Prompts are added BEFORE the position embeddings are applied (so they
+    receive position embeddings as part of the combined sequence).
+    """
+    def __init__(self, config: SiglipTextConfig):
+        super().__init__(config)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        learnable_prompts: torch.Tensor = None,
+        learning_method: str = None,
+    ) -> torch.Tensor:
+        seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+        max_position_embedding = self.position_embedding.weight.shape[0]
+
+        if seq_length > max_position_embedding:
+            raise ValueError(
+                f"Sequence length must be less than max_position_embeddings (got `sequence length`: "
+                f"{seq_length} and max_position_embeddings: {max_position_embedding}"
+            )
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.token_embedding(input_ids) # [B, L, D]
+
+        batch_size = inputs_embeds.shape[0]
+        device = inputs_embeds.device
+
+        if learnable_prompts is not None:
+            if learning_method == 'concat':
+                prompts = learnable_prompts.unsqueeze(0).expand(batch_size, -1, -1).to(device) # [B, P, D]
+                inputs_embeds = torch.cat([prompts, inputs_embeds], dim=1) # [B, P+L, D]
+            elif learning_method == 'sumate':
+                prompt_len = learnable_prompts.size(0)
+                inputs_embeds[:, :prompt_len, :] += learnable_prompts.unsqueeze(0)
+            elif learning_method == 'entire_learnable':
+                inputs_embeds = learnable_prompts.unsqueeze(0).expand(batch_size, -1, -1)
+        
+        current_len = inputs_embeds.size(1)
+        if current_len > seq_length:
+            inputs_embeds = inputs_embeds[:, :seq_length, :]
+        elif current_len < seq_length: # must check this case if happens or not
+            pad_len = seq_length - current_len
+            pad_embed = torch.zeros((batch_size, pad_len, inputs_embeds.size(2)))
+            inputs_embeds = torch.cat([inputs_embeds, pad_embed], dim=1)
+
+        position_embeddings = self.position_embedding(position_ids)
+        embeddings = inputs_embeds + position_embeddings
+
+        return embeddings
+    
+
+class SiglipTextTransformerWithPromptLearning(SiglipTextTransformer):
+    def __init__(self, config: SiglipTextConfig):
+        super().__init__(config)
+        self.embeddings = SiglipTextEmbeddingsWithPromptLearning(config)
+    
+    @can_return_tuple
+    # @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        learnable_prompts: torch.Tensor = None,
+        learning_method: str = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        if input_ids is None:
+            raise ValueError("You have to specify input_ids")
+
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids,
+                                        learnable_prompts=learnable_prompts, learning_method=learning_method)
+
+        # note: SigLIP's text model does not use a causal mask, unlike the original CLIP model.
+        # expand attention_mask
+        uses_flash_attention = "flash" in self.config._attn_implementation
+        if uses_flash_attention:
+            attention_mask = None
+        elif attention_mask is not None and not uses_flash_attention:
+            # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
+            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
+
+        encoder_outputs: BaseModelOutput = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+
+        last_hidden_state = encoder_outputs.last_hidden_state
+        last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+        # The model uses the last token's hidden state, which may be padding.
+        pooled_output = last_hidden_state[:, -1, :]
+        pooled_output = self.head(pooled_output)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+        )
+
+
+class SiglipTextModelWithPromptLearning(SiglipPreTrainedModel):
+    config: SiglipTextConfig
+    input_modalities = ("text",)
+
+    def __init__(self, config: SiglipTextConfig):
+        super().__init__(config)
+        self.text_model = SiglipTextTransformerWithPromptLearning(config)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.text_model.embeddings.token_embedding
+
+    def set_input_embeddings(self, value):
+        self.text_model.embeddings.token_embedding = value
+
+    @check_model_inputs(tie_last_hidden_states=False)
+    # @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        learnable_prompts: torch.Tensor = None,
+        learning_method: str = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import AutoTokenizer, SiglipTextModel
+
+        >>> model = SiglipTextModel.from_pretrained("google/siglip-base-patch16-224")
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224")
+
+        >>> # important: make sure to set padding="max_length" as that's how the model was trained
+        >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding="max_length", return_tensors="pt")
+
+        >>> outputs = model(**inputs)
+        >>> last_hidden_state = outputs.last_hidden_state
+        >>> pooled_output = outputs.pooler_output  # pooled (EOS token) states
+        ```"""
+
+        return self.text_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            learnable_prompts=learnable_prompts,
+            learning_method=learning_method,
+            **kwargs,
+        )

Tipsomaly/model/tips/__init__.py

@@ -0,0 +1,35 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# All software is licensed under the Apache License, Version 2.0 (Apache 2.0);
+# you may not use this file except in compliance with the Apache 2.0 license.
+# You may obtain a copy of the Apache 2.0 license at:
+# https://www.apache.org/licenses/LICENSE-2.0
+
+# All other materials are licensed under the Creative Commons Attribution 4.0
+# International License (CC-BY). You may obtain a copy of the CC-BY license at:
+# https://creativecommons.org/licenses/by/4.0/legalcode
+
+# Unless required by applicable law or agreed to in writing, all software and
+# materials distributed here under the Apache 2.0 or CC-BY licenses are
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+# either express or implied. See the licenses for the specific language
+# governing permissions and limitations under those licenses.
+
+# This is not an official Google product.
+"""Import all files."""
+# from . import text_encoder
+# from . import image_encoder
+from . import load_model

Tipsomaly/model/tips/checkpoints/checkpoint.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import os
+import subprocess
+from pathlib import Path
+from typing import Dict, List
+
+# ---- Model registry ----
+# Short unique IDs -> full checkpoint base name on GCS
+# (all files are under: https://storage.googleapis.com/tips_data/v1_0/checkpoints/pytorch/)
+MODEL_REGISTRY: Dict[str, str] = {
+    # id     # full checkpoint basename
+    "s14h":  "tips_oss_s14_highres_distilled",
+    "b14h":  "tips_oss_b14_highres_distilled",
+    "l14h":  "tips_oss_l14_highres_distilled",
+    "so4h":  "tips_oss_so400m14_highres_largetext_distilled",
+    "g14l":  "tips_oss_g14_lowres",
+    "g14h":  "tips_oss_g14_highres",
+}
+
+TOKENIZER_FILENAME = "tokenizer.model"
+TOKENIZER_URL = "https://storage.googleapis.com/tips_data/v1_0/checkpoints/tokenizer.model"
+
+BASE_URL = "https://storage.googleapis.com/tips_data/v1_0/checkpoints/pytorch"
+
+def _require_wget() -> None:
+    from shutil import which
+    if which("wget") is None:
+        raise EnvironmentError(
+            "wget is required but was not found on PATH. "
+            "Please install wget or add it to your PATH."
+        )
+
+def _wget(url: str, dest: Path) -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    # -q for quiet except errors; --show-progress if attached to TTY would be nice,
+    # but we keep it simple and quiet here.
+    result = subprocess.run(
+        ["wget", "-q", "-O", str(dest), url],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    if result.returncode != 0:
+        # Clean up partial file if any
+        if dest.exists() and dest.stat().st_size == 0:
+            try:
+                dest.unlink()
+            except OSError:
+                pass
+        raise RuntimeError(
+            f"Failed to download {url} -> {dest}\n"
+            f"wget stderr:\n{result.stderr.strip()}"
+        )
+
+def _model_files_for_basename(base: str) -> List[str]:
+    return [f"{base}_vision.npz", f"{base}_text.npz"]
+
+def list_models() -> Dict[str, str]:
+    return dict(MODEL_REGISTRY)
+
+def ensure_model_files(model_id: str, save_dir: str) -> Dict[str, str]:
+    if model_id not in MODEL_REGISTRY:
+        raise KeyError(
+            f"Unknown model_id '{model_id}'. "
+            f"Valid options: {', '.join(sorted(MODEL_REGISTRY.keys()))}"
+        )
+
+    _require_wget()
+
+    save_path = Path(save_dir).expanduser().resolve()
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    # 1) Tokenizer
+    tokenizer_path = save_path / TOKENIZER_FILENAME
+    if not tokenizer_path.exists():
+        _wget(TOKENIZER_URL, tokenizer_path)
+
+    # 2) Model files
+    base = MODEL_REGISTRY[model_id]
+    required_files = _model_files_for_basename(base)
+
+    out_paths = {"tokenizer": str(tokenizer_path)}
+
+    for fname in required_files:
+        local_path = save_path / fname
+        if not local_path.exists():
+            # Compose the correct URL for this file
+            url = f"{BASE_URL}/{fname}"
+            _wget(url, local_path)
+
+        # record
+        if fname.endswith("_vision.npz"):
+            out_paths["vision"] = str(local_path)
+        elif fname.endswith("_text.npz"):
+            out_paths["text"] = str(local_path)
+
+    return out_paths

Tipsomaly/model/tips/checkpoints/download_checkpoints.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+# The model weights can be found in https://console.cloud.google.com/storage/browser/tips_data
+ALL_CHECKPOINTS=(
+  "tips_oss_s14_highres_distilled"
+  "tips_oss_b14_highres_distilled"
+  "tips_oss_l14_highres_distilled"
+  "tips_oss_so400m14_highres_largetext_distilled"
+  "tips_oss_g14_lowres"
+  "tips_oss_g14_highres"
+)
+
+echo "Downloading the tokenizer."
+wget https://storage.googleapis.com/tips_data/v1_0/checkpoints/tokenizer.model
+
+for CHECKPOINT in "${ALL_CHECKPOINTS[@]}"; do
+  echo "Downloading ${CHECKPOINT} (vision encoder weights)"
+  wget https://storage.googleapis.com/tips_data/v1_0/checkpoints/pytorch/${CHECKPOINT}_vision.npz
+  echo "Downloading ${CHECKPOINT} (text encoder weights)"
+  wget https://storage.googleapis.com/tips_data/v1_0/checkpoints/pytorch/${CHECKPOINT}_text.npz
+done

Tipsomaly/model/tips/image_encoder.py

@@ -0,0 +1,1002 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Vision encoder implementation in PyTorch."""
+
+import functools
+import math
+import os
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+import warnings
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+
+class Mlp(nn.Module):
+  """Transformer MLP, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      in_features: int,
+      hidden_features: Optional[int] = None,
+      out_features: Optional[int] = None,
+      act_layer: Callable[..., nn.Module] = nn.GELU,
+      drop: float = 0.0,
+      bias: bool = True,
+  ) -> None:
+    super().__init__()
+    out_features = out_features or in_features
+    hidden_features = hidden_features or in_features
+    self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+    self.act = act_layer()
+    self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+    self.drop = nn.Dropout(drop)
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    x = self.fc1(x)
+    x = self.act(x)
+    x = self.drop(x)
+    x = self.fc2(x)
+    x = self.drop(x)
+    return x
+
+
+def make_2tuple(x):
+  if isinstance(x, tuple):
+    assert len(x) == 2
+    return x
+
+  assert isinstance(x, int)
+  return (x, x)
+
+
+class PatchEmbed(nn.Module):
+  """2D image to patch embedding: (B,C,H,W) -> (B,N,D)."""
+
+  def __init__(
+      self,
+      img_size: Union[int, Tuple[int, int]] = 224,
+      patch_size: Union[int, Tuple[int, int]] = 16,
+      in_chans: int = 3,
+      embed_dim: int = 768,
+      norm_layer: Optional[Callable] = None,  # pylint: disable=g-bare-generic
+      flatten_embedding: bool = True,
+  ) -> None:
+    super().__init__()
+
+    image_hw = make_2tuple(img_size)
+    patch_hw = make_2tuple(patch_size)
+    patch_grid_size = (
+        image_hw[0] // patch_hw[0],
+        image_hw[1] // patch_hw[1],
+    )
+
+    self.img_size = image_hw
+    self.patch_size = patch_hw
+    self.patches_resolution = patch_grid_size
+    self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+    self.in_chans = in_chans
+    self.embed_dim = embed_dim
+
+    self.flatten_embedding = flatten_embedding
+
+    self.proj = nn.Conv2d(
+        in_chans, embed_dim, kernel_size=patch_hw, stride=patch_hw
+    )
+    self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    _, _, h, w = x.shape
+    patch_h, patch_w = self.patch_size
+
+    assert (
+        h % patch_h == 0
+    ), f"Input image height {h} is not a multiple of patch height {patch_h}"
+    assert (
+        w % patch_w == 0
+    ), f"Input image width {w} is not a multiple of patch width: {patch_w}"
+
+    x = self.proj(x)  # B C H W
+    h, w = x.size(2), x.size(3)
+    x = x.flatten(2).transpose(1, 2)  # B HW C
+    x = self.norm(x)
+    if not self.flatten_embedding:
+      x = x.reshape(-1, h, w, self.embed_dim)  # B H W C
+    return x
+
+  def flops(self) -> float:
+    ho, wo = self.patches_resolution
+    flops = (
+        ho
+        * wo
+        * self.embed_dim
+        * self.in_chans
+        * (self.patch_size[0] * self.patch_size[1])
+    )
+    if self.norm is not None:
+      flops += ho * wo * self.embed_dim
+    return flops
+
+
+class SwiGLUFFN(nn.Module):
+  """SwiGLU FFN layer, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      in_features: int,
+      hidden_features: Optional[int] = None,
+      out_features: Optional[int] = None,
+      act_layer: Callable[..., nn.Module] = None,
+      drop: float = 0.0,
+      bias: bool = True,
+  ) -> None:
+    super().__init__()
+    out_features = out_features or in_features
+    hidden_features = hidden_features or in_features
+    self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+    self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    x12 = self.w12(x)
+    x1, x2 = x12.chunk(2, dim=-1)
+    hidden = F.silu(x1) * x2
+    return self.w3(hidden)
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+  if XFORMERS_ENABLED:
+    from xformers.ops import SwiGLU, memory_efficient_attention, unbind, fmha, scaled_index_add, index_select_cat  # pylint: disable=g-multiple-import, g-import-not-at-top
+
+    XFORMERS_AVAILABLE = True
+    warnings.warn("xFormers is available (SwiGLU)")
+  else:
+    warnings.warn("xFormers is disabled (SwiGLU)")
+    raise ImportError
+except ImportError:
+  SwiGLU = SwiGLUFFN
+  XFORMERS_AVAILABLE = False
+
+  warnings.warn("xFormers is not available (SwiGLU)")
+
+
+class SwiGLUFFNFused(SwiGLU):
+  """SwiGLU FFN layer, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      in_features: int,
+      hidden_features: Optional[int] = None,
+      out_features: Optional[int] = None,
+      act_layer: Callable[..., nn.Module] = None,  # pylint: disable=unused-argument
+      drop: float = 0.0,  # pylint: disable=unused-argument
+      bias: bool = True,
+  ) -> None:
+    out_features = out_features or in_features
+    hidden_features = hidden_features or in_features
+    hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+    super().__init__(
+        in_features=in_features,
+        hidden_features=hidden_features,
+        out_features=out_features,
+        bias=bias,
+    )
+
+
+class Attention(nn.Module):
+  """Attention layer, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      dim: int,
+      num_heads: int = 8,
+      qkv_bias: bool = False,
+      proj_bias: bool = True,
+      attn_drop: float = 0.0,
+      proj_drop: float = 0.0,
+  ) -> None:
+    super().__init__()
+    self.num_heads = num_heads
+    head_dim = dim // num_heads
+    self.scale = head_dim**-0.5
+
+    self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+    self.attn_drop = nn.Dropout(attn_drop)
+    self.proj = nn.Linear(dim, dim, bias=proj_bias)
+    self.proj_drop = nn.Dropout(proj_drop)
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    b_dim, n_dim, c_dim = x.shape
+    qkv = (
+        self.qkv(x)
+        .reshape(b_dim, n_dim, 3, self.num_heads, c_dim // self.num_heads)
+        .permute(2, 0, 3, 1, 4)
+    )
+
+    q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+    attn = q @ k.transpose(-2, -1)
+
+    attn = attn.softmax(dim=-1)
+    attn = self.attn_drop(attn)
+
+    x = (attn @ v).transpose(1, 2).reshape(b_dim, n_dim, c_dim)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+
+class MemEffAttention(Attention):
+  """Memory Efficient Attention layer, following DINOv2 implementation."""
+
+  def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+    if not XFORMERS_AVAILABLE:
+      if attn_bias is not None:
+        raise AssertionError("xFormers is required for using nested tensors")
+      return super().forward(x)
+
+    b_dim, n_dim, c_dim = x.shape
+    qkv = self.qkv(x).reshape(
+        b_dim, n_dim, 3, self.num_heads, c_dim // self.num_heads
+    )
+
+    q, k, v = unbind(qkv, 2)
+
+    x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+    x = x.reshape([b_dim, n_dim, c_dim])
+
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+
+class LayerScale(nn.Module):
+  """Layer scale, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      dim: int,
+      init_values: Union[float, torch.Tensor] = 1e-5,
+      inplace: bool = False,
+  ) -> None:
+    super().__init__()
+    self.inplace = inplace
+    self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+def drop_path_impl(x, drop_prob: float = 0.0, training: bool = False):
+  if drop_prob == 0.0 or not training:
+    return x
+  keep_prob = 1 - drop_prob
+  shape = (x.shape[0],) + (1,) * (
+      x.ndim - 1
+  )  # work with diff dim tensors, not just 2D ConvNets
+  random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+  if keep_prob > 0.0:
+    random_tensor.div_(keep_prob)
+  output = x * random_tensor
+  return output
+
+
+class DropPath(nn.Module):
+  """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+  def __init__(self, drop_prob=None):
+    super(DropPath, self).__init__()
+    self.drop_prob = drop_prob
+
+  def forward(self, x):
+    return drop_path_impl(x, self.drop_prob, self.training)
+
+
+class Block(nn.Module):
+  """Transformer Block Implementation, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      dim: int,
+      num_heads: int,
+      mlp_ratio: float = 4.0,
+      qkv_bias: bool = False,
+      proj_bias: bool = True,
+      ffn_bias: bool = True,
+      drop: float = 0.0,
+      attn_drop: float = 0.0,
+      init_values=None,
+      drop_path: float = 0.0,
+      act_layer: Callable[..., nn.Module] = nn.GELU,
+      norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+      attn_class: Callable[..., nn.Module] = Attention,
+      ffn_layer: Callable[..., nn.Module] = Mlp,
+  ) -> None:
+    super().__init__()
+    self.norm1 = norm_layer(dim)
+    self.attn = attn_class(
+        dim,
+        num_heads=num_heads,
+        qkv_bias=qkv_bias,
+        proj_bias=proj_bias,
+        attn_drop=attn_drop,
+        proj_drop=drop,
+    )
+    self.ls1 = (
+        LayerScale(dim, init_values=init_values)
+        if init_values
+        else nn.Identity()
+    )
+    self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    self.norm2 = norm_layer(dim)
+    mlp_hidden_dim = int(dim * mlp_ratio)
+    self.mlp = ffn_layer(
+        in_features=dim,
+        hidden_features=mlp_hidden_dim,
+        act_layer=act_layer,
+        drop=drop,
+        bias=ffn_bias,
+    )
+    self.ls2 = (
+        LayerScale(dim, init_values=init_values)
+        if init_values
+        else nn.Identity()
+    )
+    self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    self.sample_drop_ratio = drop_path
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def attn_residual_func(x: torch.Tensor) -> torch.Tensor:
+      return self.ls1(self.attn(self.norm1(x)))
+
+    def ffn_residual_func(x: torch.Tensor) -> torch.Tensor:
+      return self.ls2(self.mlp(self.norm2(x)))
+
+    if self.training and self.sample_drop_ratio > 0.1:
+      # the overhead is compensated only for a drop path rate larger than 0.1
+      x = drop_add_residual_stochastic_depth(
+          x,
+          residual_func=attn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+      )
+      x = drop_add_residual_stochastic_depth(
+          x,
+          residual_func=ffn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+      )
+    elif self.training and self.sample_drop_ratio > 0.0:
+      x = x + self.drop_path1(attn_residual_func(x))
+      x = x + self.drop_path1(ffn_residual_func(x))
+    else:
+      x = x + attn_residual_func(x)
+      x = x + ffn_residual_func(x)
+    return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: torch.Tensor,
+    residual_func: Callable[[torch.Tensor], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> torch.Tensor:
+  """This function is taken from the original implementation in DINOv2 to implement stochastic depth in the image encoder."""
+  # 1) extract subset using permutation
+  b, _, _ = x.shape
+  sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+  brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+  x_subset = x[brange]
+
+  # 2) apply residual_func to get residual
+  residual = residual_func(x_subset)
+
+  x_flat = x.flatten(1)
+  residual = residual.flatten(1)
+
+  residual_scale_factor = b / sample_subset_size
+
+  # 3) add the residual
+  x_plus_residual = torch.index_add(
+      x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor
+  )
+  return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+  b, _, _ = x.shape
+  sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+  brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+  residual_scale_factor = b / sample_subset_size
+  return brange, residual_scale_factor
+
+
+def add_residual(
+    x, brange, residual, residual_scale_factor, scaling_vector=None
+):
+  """Implement residual addition in the image encoder."""
+  if scaling_vector is None:
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+    x_plus_residual = torch.index_add(
+        x_flat,
+        0,
+        brange,
+        residual.to(dtype=x.dtype),
+        alpha=residual_scale_factor,
+    )
+  else:
+    x_plus_residual = scaled_index_add(
+        x,
+        brange,
+        residual.to(dtype=x.dtype),
+        scaling=scaling_vector,
+        alpha=residual_scale_factor,
+    )
+  return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}  # pylint: disable=g-bare-generic
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+  """this will perform the index select, cat the tensors, and provide the attn_bias from cache."""
+  batch_sizes = (
+      [b.shape[0] for b in branges]
+      if branges is not None
+      else [x.shape[0] for x in x_list]
+  )
+  all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+  if all_shapes not in attn_bias_cache.keys():
+    seqlens = []
+    for b, x in zip(batch_sizes, x_list):
+      for _ in range(b):
+        seqlens.append(x.shape[1])
+    attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+    attn_bias._batch_sizes = batch_sizes  # pylint: disable=protected-access
+    attn_bias_cache[all_shapes] = attn_bias
+
+  if branges is not None:
+    cat_tensors = index_select_cat(
+        [x.flatten(1) for x in x_list], branges
+    ).view(1, -1, x_list[0].shape[-1])
+  else:
+    tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+    cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+  return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[torch.Tensor],
+    residual_func: Callable[[torch.Tensor, Any], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> torch.Tensor:
+  """Add residual to a list of tensors."""
+  # 1) generate random set of indices for dropping samples in the batch.
+  branges_scales = [
+      get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list
+  ]
+  branges = [s[0] for s in branges_scales]
+  residual_scale_factors = [s[1] for s in branges_scales]
+
+  # 2) get attention bias and index+concat the tensors.
+  attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+  # 3) apply residual_func to get residual, and split the result.
+  residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+  outputs = []
+  for x, brange, residual, residual_scale_factor in zip(
+      x_list, branges, residual_list, residual_scale_factors
+  ):
+    outputs.append(
+        add_residual(
+            x, brange, residual, residual_scale_factor, scaling_vector
+        ).view_as(x)
+    )
+  return outputs
+
+
+class NestedTensorBlock(Block):
+  """Nested tensor block implementation."""
+
+  def forward_nested(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
+    """x_list contains a list of tensors to nest together and run."""
+    assert isinstance(self.attn, MemEffAttention)
+
+    if self.training and self.sample_drop_ratio > 0.0:
+
+      def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+      def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        del attn_bias
+        return self.mlp(self.norm2(x))
+
+      x_list = drop_add_residual_stochastic_depth_list(
+          x_list,
+          residual_func=attn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+          scaling_vector=self.ls1.gamma
+          if isinstance(self.ls1, LayerScale)
+          else None,
+      )
+      x_list = drop_add_residual_stochastic_depth_list(
+          x_list,
+          residual_func=ffn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+          scaling_vector=self.ls2.gamma
+          if isinstance(self.ls1, LayerScale)
+          else None,
+      )
+      return x_list
+    else:
+
+      def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+      def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        del attn_bias
+        return self.ls2(self.mlp(self.norm2(x)))
+
+      attn_bias, x = get_attn_bias_and_cat(x_list)
+      x = x + attn_residual_func(x, attn_bias=attn_bias)
+      x = x + ffn_residual_func(x)
+      return attn_bias.split(x)
+
+  def forward(self, x):
+    if isinstance(x, torch.Tensor):
+      return super().forward(x)
+    elif isinstance(x, list):
+      if not XFORMERS_AVAILABLE:
+        raise AssertionError("xFormers is required for using nested tensors")
+      return self.forward_nested(x)
+    else:
+      raise AssertionError
+
+
+def named_apply(
+    fn: Callable,  # pylint: disable=g-bare-generic
+    module: nn.Module,
+    name="",
+    depth_first=True,
+    include_root=False,
+) -> nn.Module:
+  """Apply a function to a module and its children."""
+  if not depth_first and include_root:
+    fn(module=module, name=name)
+  for child_name, child_module in module.named_children():
+    child_name = ".".join((name, child_name)) if name else child_name
+    named_apply(
+        fn=fn,
+        module=child_module,
+        name=child_name,
+        depth_first=depth_first,
+        include_root=True,
+    )
+  if depth_first and include_root:
+    fn(module=module, name=name)
+  return module
+
+
+class BlockChunk(nn.ModuleList):
+
+  def forward(self, x):
+    for b in self:
+      x = b(x)
+    return x
+
+
+class VisionTransformer(nn.Module):
+  """Vision Transformer implementation."""
+
+  def __init__(
+      self,
+      img_size=224,
+      patch_size=16,
+      in_chans=3,
+      embed_dim=768,
+      depth=12,
+      num_heads=12,
+      mlp_ratio=4.0,
+      qkv_bias=True,
+      ffn_bias=True,
+      proj_bias=True,
+      drop_path_rate=0.0,
+      drop_path_uniform=False,
+      init_values=None,  # for layerscale: None or 0 => no layerscale
+      embed_layer=PatchEmbed,
+      act_layer=nn.GELU,
+      block_fn=Block,
+      ffn_layer="mlp",
+      block_chunks=1,
+      num_register_tokens=0,
+      interpolate_antialias=False,
+      interpolate_offset=0.1,
+  ):
+    """Defines the Vision Transformer model.
+
+    Args:
+      img_size (int, tuple): input image size
+      patch_size (int, tuple): patch size
+      in_chans (int): number of input channels
+      embed_dim (int): embedding dimension
+      depth (int): depth of transformer
+      num_heads (int): number of attention heads
+      mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+      qkv_bias (bool): enable bias for qkv if True
+      ffn_bias (bool): enable bias for ffn if True
+      proj_bias (bool): enable bias for proj in attn if True
+      drop_path_rate (float): stochastic depth rate
+      drop_path_uniform (bool): apply uniform drop rate across blocks
+      init_values (float): layer-scale init values
+      embed_layer (nn.Module): patch embedding layer
+      act_layer (nn.Module): MLP activation layer
+      block_fn (nn.Module): transformer block class
+      ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+      block_chunks: (int) split block sequence into block_chunks units for FSDP
+        wrap
+      num_register_tokens: (int) number of extra cls tokens (so-called
+        "registers")
+      interpolate_antialias: (str) flag to apply anti-aliasing when
+        interpolating positional embeddings
+      interpolate_offset: (float) work-around offset to apply when interpolating
+        positional embeddings
+    """
+    super().__init__()
+    norm_layer = functools.partial(nn.LayerNorm, eps=1e-6)
+
+    self.num_features = self.embed_dim = (
+        embed_dim  # num_features for consistency with other models
+    )
+    self.num_tokens = 1
+    self.n_blocks = depth
+    self.num_heads = num_heads
+    self.patch_size = patch_size
+    self.num_register_tokens = num_register_tokens
+    self.interpolate_antialias = interpolate_antialias
+    self.interpolate_offset = interpolate_offset
+
+    self.patch_embed = embed_layer(
+        img_size=img_size,
+        patch_size=patch_size,
+        in_chans=in_chans,
+        embed_dim=embed_dim,
+    )
+    num_patches = self.patch_embed.num_patches
+
+    self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+    self.pos_embed = nn.Parameter(
+        torch.zeros(1, num_patches + self.num_tokens, embed_dim)
+    )
+    assert num_register_tokens >= 0
+    self.register_tokens = (
+        nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim))
+        if num_register_tokens
+        else None
+    )
+
+    if drop_path_uniform:
+      dpr = [drop_path_rate] * depth
+    else:
+      dpr = [
+          x.item() for x in torch.linspace(0, drop_path_rate, depth)
+      ]  # stochastic depth decay rule
+
+    if ffn_layer == "mlp":
+      ffn_layer = Mlp
+    elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+      ffn_layer = SwiGLUFFNFused
+    else:
+      raise NotImplementedError
+
+    blocks_list = [
+        block_fn(
+            dim=embed_dim,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            ffn_bias=ffn_bias,
+            drop_path=dpr[i],
+            norm_layer=norm_layer,
+            act_layer=act_layer,
+            ffn_layer=ffn_layer,
+            init_values=init_values,
+        )
+        for i in range(depth)
+    ]
+    if block_chunks > 0:
+      self.chunked_blocks = True
+      chunked_blocks = []
+      chunksize = depth // block_chunks
+      for i in range(0, depth, chunksize):
+        # this is to keep the block index consistent if we chunk the block list
+        chunked_blocks.append(
+            [nn.Identity()] * i + blocks_list[i : i + chunksize]
+        )
+      self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+    else:
+      self.chunked_blocks = False
+      self.blocks = nn.ModuleList(blocks_list)
+
+    self.norm = norm_layer(embed_dim)
+    self.head = nn.Identity()
+
+    self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+    self.init_weights()
+
+  def init_weights(self):
+    nn.init.trunc_normal_(self.pos_embed, std=0.02)
+    nn.init.normal_(self.cls_token, std=1e-6)
+    if self.register_tokens is not None:
+      nn.init.normal_(self.register_tokens, std=1e-6)
+    named_apply(init_weights_vit_timm, self)
+
+  def interpolate_pos_encoding(self, x, w, h):
+    previous_dtype = x.dtype
+    npatch = x.shape[1] - 1
+    num_patches = self.pos_embed.shape[1] - 1
+    if npatch == num_patches and w == h:
+      return self.pos_embed
+    pos_embed = self.pos_embed.float()
+    class_pos_embed = pos_embed[:, 0]
+    patch_pos_embed = pos_embed[:, 1:]
+    dim = x.shape[-1]
+    w0 = w // self.patch_size
+    h0 = h // self.patch_size
+    num_patches_dim = int(
+        math.sqrt(num_patches)
+    )  # Recover the number of patches in each dimension
+    assert num_patches == num_patches_dim * num_patches_dim
+    kwargs = {}
+    if self.interpolate_offset:
+      sx = float(w0 + self.interpolate_offset) / num_patches_dim
+      sy = float(h0 + self.interpolate_offset) / num_patches_dim
+      kwargs["scale_factor"] = (sx, sy)
+    else:
+      # Simply specify an output size instead of a scale factor
+      kwargs["size"] = (w0, h0)
+    patch_pos_embed = nn.functional.interpolate(
+        patch_pos_embed.reshape(
+            1, num_patches_dim, num_patches_dim, dim
+        ).permute(0, 3, 1, 2),
+        mode="bilinear",
+        antialias=self.interpolate_antialias,
+        **kwargs,
+    )
+    assert (w0, h0) == patch_pos_embed.shape[-2:]
+    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+    return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(
+        previous_dtype
+    )
+
+  def prepare_tokens_with_masks(self, x, masks=None):
+    _, _, w, h = x.shape
+    x = self.patch_embed(x)
+    if masks is not None:
+      x = torch.where(
+          masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x
+      )
+
+    x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+    x = x + self.interpolate_pos_encoding(x, w, h)
+
+    if self.register_tokens is not None:
+      x = torch.cat(
+          (
+              x[:, :1],
+              self.register_tokens.expand(x.shape[0], -1, -1),
+              x[:, 1:],
+          ),
+          dim=1,
+      )
+
+    return x
+
+  def forward_features_list(self, x_list, masks_list):
+    x = [
+        self.prepare_tokens_with_masks(x, masks)
+        for x, masks in zip(x_list, masks_list)
+    ]
+    for blk in self.blocks:
+      x = blk(x)
+
+    all_x = x
+    output = []
+    for x, masks in zip(all_x, masks_list):
+      x_norm = self.norm(x)
+      output.append({
+          "x_norm_1st_clstoken": x_norm[:, :1],
+          "x_norm_2nd_clstoken": x_norm[:, 1 : self.num_register_tokens + 1],
+          "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+          "x_prenorm": x,
+          "masks": masks,
+      })
+    return output
+
+  def forward_features(self, x, masks=None):
+    if isinstance(x, list):
+      return self.forward_features_list(x, masks)
+
+    x = self.prepare_tokens_with_masks(x, masks)
+
+    for blk in self.blocks:
+      x = blk(x)
+
+    x_norm = self.norm(x)
+    return {
+        "x_norm_1st_clstoken": x_norm[:, :1],
+        "x_norm_2nd_clstoken": x_norm[:, 1 : self.num_register_tokens + 1],
+        "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+        "x_prenorm": x,
+        "masks": masks,
+    }
+
+  def _get_intermediate_layers_not_chunked(self, x, n=1):
+    x = self.prepare_tokens_with_masks(x)
+    # If n is an int, take the n last blocks. If it's a list, take them
+    output, total_block_len = [], len(self.blocks)
+    blocks_to_take = (
+        range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+    )
+    for i, blk in enumerate(self.blocks):
+      x = blk(x)
+      if i in blocks_to_take:
+        output.append(x)
+    assert len(output) == len(
+        blocks_to_take
+    ), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+    return output
+
+  def _get_intermediate_layers_chunked(self, x, n=1):
+    x = self.prepare_tokens_with_masks(x)
+    output, i, total_block_len = [], 0, len(self.blocks[-1])
+    # If n is an int, take the n last blocks. If it's a list, take them
+    blocks_to_take = (
+        range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+    )
+    for block_chunk in self.blocks:
+      for blk in block_chunk[i:]:  # Passing the nn.Identity()
+        x = blk(x)
+        if i in blocks_to_take:
+          output.append(x)
+        i += 1
+    assert len(output) == len(
+        blocks_to_take
+    ), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+    return output
+
+  def get_intermediate_layers(
+      self,
+      x: torch.torch.Tensor,
+      n: Union[int, Sequence] = 1,  # Layers or n last layers to take  # pylint: disable=g-bare-generic
+      reshape: bool = False,
+      return_class_token: bool = False,
+      norm=True,
+  ) -> Tuple[Union[torch.torch.Tensor, Tuple[torch.torch.Tensor]]]:  # pylint: disable=g-one-element-tuple
+    if self.chunked_blocks:
+      outputs = self._get_intermediate_layers_chunked(x, n)
+    else:
+      outputs = self._get_intermediate_layers_not_chunked(x, n)
+    if norm:
+      outputs = [self.norm(out) for out in outputs]
+    class_tokens = [out[:, 0] for out in outputs]
+    outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
+    if reshape:
+      batch_size, _, w, h = x.shape
+      outputs = [
+          out.reshape(
+              batch_size, w // self.patch_size, h // self.patch_size, -1
+          )
+          .permute(0, 3, 1, 2)
+          .contiguous()
+          for out in outputs
+      ]
+    if return_class_token:
+      return tuple(zip(outputs, class_tokens))
+    return tuple(outputs)
+
+  def forward(self, *args, is_training=False, **kwargs):
+    ret = self.forward_features(*args, **kwargs)
+    if is_training:
+      return ret
+    else:
+      return self.head(ret["x_norm_1st_clstoken"]), self.head(
+          ret["x_norm_2nd_clstoken"]
+      ), ret["x_norm_patchtokens"]
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):  # pylint: disable=unused-argument
+  """ViT weight initialization, original timm impl (for reproducibility)."""
+  if isinstance(module, nn.Linear):
+    nn.init.trunc_normal_(module.weight, std=0.02)
+    if module.bias is not None:
+      nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=384,
+      depth=12,
+      num_heads=6,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_base(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=768,
+      depth=12,
+      num_heads=12,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_large(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=1024,
+      depth=24,
+      num_heads=16,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_so400m(patch_size=14, **kwargs):
+  """SoViT 400M model (https://arxiv.org/abs/2305.13035)."""
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=1152,
+      depth=27,
+      num_heads=16,
+      mlp_ratio=4304 / 1152,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_giant2(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=1536,
+      depth=40,
+      num_heads=24,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model

Tipsomaly/model/tips/load_model.py

@@ -0,0 +1,112 @@
+import io
+import numpy as np
+import torch
+import os
+
+from .text_encoder import TextEncoder, Tokenizer
+# from .image_encoder import vit_small, vit_base, vit_large, vit_so400m, vit_giant2
+from .image_encoder import Block, VisionTransformer, MemEffAttention
+import functools
+
+from .checkpoints import checkpoint
+
+MAX_LEN = 64
+VOCAB_SIZE = 32000
+PATCH_SIZE = 14
+
+vision_models = {
+    'S':      {'embed_dim': 384,  'depth': 12, 'num_heads': 6,  'mlp_ratio': 4.0},
+    'B':      {'embed_dim': 768,  'depth': 12, 'num_heads': 12, 'mlp_ratio': 4.0},
+    'L':      {'embed_dim': 1024, 'depth': 24, 'num_heads': 16, 'mlp_ratio': 4.0},
+    'So400m': {'embed_dim': 1152, 'depth': 27, 'num_heads': 16, 'mlp_ratio': 4304/1152},
+    'G':      {'embed_dim': 1536, 'depth': 40, 'num_heads': 24, 'mlp_ratio': 4.0},
+}
+def build_vision_encoder(cfg: dict, *, img_size, patch_size, ffn_layer):
+    return VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=cfg['embed_dim'],
+        depth=cfg['depth'],
+        num_heads=cfg['num_heads'],
+        mlp_ratio=cfg['mlp_ratio'],
+        block_fn=functools.partial(Block, attn_class=MemEffAttention),
+        num_register_tokens=1,
+        img_size=img_size,
+        ffn_layer=ffn_layer,
+        block_chunks=0,
+        init_values=1.0,
+        interpolate_antialias=True,
+        interpolate_offset=0.0,
+    )
+    
+def load_image_encoder(model_weights_path, model_variant, is_low_res, patch_size = 14):
+    img_size = 224 if is_low_res else 448
+    ffn_layer = 'swiglu' if model_variant == 'G' else 'mlp'
+    cfg = vision_models[model_variant]
+
+    checkpoint_np = dict(np.load(model_weights_path, allow_pickle=False))
+    checkpoint = {k: torch.tensor(v) for k, v in checkpoint_np.items()}
+
+    with torch.no_grad():
+        model = build_vision_encoder(cfg, img_size=img_size, patch_size=patch_size, ffn_layer=ffn_layer)
+        missing, unexpected = model.load_state_dict(checkpoint, strict=False)
+        # Optional: sanity logs
+        if missing:
+            print(f"[vision:{model_variant}] Missing keys: {len(missing)} (e.g. {missing[:3]})")
+        if unexpected:
+            print(f"[vision:{model_variant}] Unexpected keys: {len(unexpected)} (e.g. {unexpected[:3]})")
+
+    return model
+  
+text_models = {
+    'S': {'hidden_size': 384, 'mlp_dim': 1536, 'num_heads': 6, 'num_layers': 12},
+    'B': {'hidden_size': 768, 'mlp_dim': 3072, 'num_heads': 12, 'num_layers': 12},
+    'L': {'hidden_size': 1024, 'mlp_dim': 4096, 'num_heads': 16, 'num_layers': 12},
+    'So400m': {'hidden_size': 1152, 'mlp_dim': 4304, 'num_heads': 16, 'num_layers': 27},
+    'G': {'hidden_size': 1536, 'mlp_dim': 6144, 'num_heads': 24, 'num_layers': 12},
+}
+def load_text_encoder(model_path, model_variant, tokenizer_path):
+    with open(model_path, 'rb') as fin:
+        inbuffer = io.BytesIO(fin.read())
+    np_weights_text = np.load(inbuffer, allow_pickle=False)
+
+    weights_text = {}
+    for key, value in np_weights_text.items():
+        weights_text[key] = torch.from_numpy(value)
+    temperature = weights_text.pop('temperature')
+
+    with torch.no_grad():
+        # Define the text model.
+        model_text = TextEncoder(
+            text_models[model_variant],
+            vocab_size=VOCAB_SIZE,
+        )
+        model_text.load_state_dict(weights_text)
+
+        tokenizer_obj = Tokenizer(tokenizer_path=tokenizer_path)
+    return model_text, tokenizer_obj, temperature
+
+
+CHECKPOINT_TO_VARIANT = {
+    "s14h": "S",
+    "b14h": "B",
+    "l14h": "L",
+    "so4h": "So400m",   # 400m = big SoViT model
+    "g14l": "G",        # Giant low-res
+    "g14h": "G",        # Giant high-res
+}
+def get_model(model_path, model_checkpoint):
+    paths = checkpoint.ensure_model_files(model_checkpoint, model_path)
+    for key, path in paths.items():
+        print(f"{key}: {path}")
+    
+    tokenizer_path = os.path.join(model_path, 'tokenizer.model')
+    image_enc_name, text_enc_name = checkpoint._model_files_for_basename(checkpoint.MODEL_REGISTRY[model_checkpoint])
+    image_encoder_path = os.path.join(model_path, image_enc_name)
+    text_encoder_path = os.path.join(model_path, text_enc_name)
+    is_low_res = model_checkpoint.endswith('l')
+    
+    model_variant = CHECKPOINT_TO_VARIANT[model_checkpoint]
+    image_encoder = load_image_encoder(image_encoder_path, model_variant, is_low_res)
+    text_encoder, tokenizer, temperature = load_text_encoder(text_encoder_path, model_variant, tokenizer_path)
+    
+    return image_encoder, text_encoder, tokenizer, temperature

Tipsomaly/model/tips/text_encoder.py

@@ -0,0 +1,519 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Text encoder implementation in PyTorch."""
+
+import typing as t
+
+# import tensorflow as tf
+# import tensorflow_text
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+import torch
+import sentencepiece as spm
+
+class Tokenizer:
+    """A simple tokenizer using SentencePiece with batch support."""
+
+    def __init__(self, tokenizer_path: str):
+        """Initializes the tokenizer."""
+        self.tokenizer = spm.SentencePieceProcessor(model_file=tokenizer_path)
+
+    def tokenize(self, input_texts, max_len=64):
+        """Tokenizes a batch of input texts and pads/clips them to max_len."""
+        if isinstance(input_texts, str):
+            input_texts = [input_texts]  # Convert single string to list
+        
+        input_texts = [text.lower() for text in input_texts]  # Lowercasing batch
+        tokenized = [self.tokenizer.encode(text) for text in input_texts]  # Tokenize batch
+        
+        # Pad or truncate to max_len
+        padded_tokens = []
+        padding_masks = []
+        for tokens in tokenized:
+            curr_len = len(tokens)
+            if curr_len > max_len:
+                tokens = tokens[:max_len]  # Truncate
+            else:
+                tokens += [0] * (max_len - curr_len)  # Pad
+            
+            padded_tokens.append(tokens)
+            padding_masks.append([0] * curr_len + [1] * (max_len - curr_len))  # Mask (1 for real tokens, 0 for padding)
+        
+        # Convert to PyTorch tensors
+        tokens_tensor = torch.tensor(padded_tokens, dtype=torch.int32)
+        padding_mask_tensor = torch.tensor(padding_masks, dtype=torch.int32)
+        
+        return tokens_tensor, padding_mask_tensor
+
+
+# class Tokenizer(object):
+#   """A simple tokenizer."""
+
+#   def __init__(self, tokenizer_path: str):
+#     """Initializes the tokenizer."""
+#     with open(tokenizer_path, 'rb') as f:
+#       model = f.read()
+#     self.tokenizer = tensorflow_text.SentencepieceTokenizer(
+#         model=model, add_eos=False, add_bos=False
+#     )
+
+#   def tokenize(self, input_text, max_len=64):
+#     tokens = self.tokenizer.tokenize(tf.strings.lower(input_text)).to_tensor()
+#     curr_len = tokens.shape[1]
+#     is_padding = tf.zeros((tokens.shape[0], max_len))
+#     if curr_len > max_len:
+#       tokens = tokens[:, :max_len]
+#     else:
+#       padding_len = max_len - curr_len
+#       tokens = tf.pad(tokens, [[0, 0], [0, padding_len]], constant_values=0)
+#       is_padding = tf.cast(tokens == 0, tf.int32)
+#     return tokens.numpy(), is_padding.numpy()
+
+
+class PositionalEmbedding(nn.Module):
+  """Generates position embedding for a given 1-d sequence.
+
+  Attributes:
+    min_timescale: Start of the geometric index. Determines the periodicity of
+      the added signal.
+    max_timescale: End of the geometric index. Determines the frequency of the
+      added signal.
+    embedding_dim: Dimension of the embedding to be generated.
+  """
+
+  min_timescale: int = 1
+  max_timescale: int = 10_000
+  embedding_dim: int = 0
+
+  def __init__(self, embedding_dim: int):
+    super().__init__()
+    self.embedding_dim = embedding_dim
+
+  def __call__(self, seq_length: int = None, position: torch.tensor = None):
+    """Generates a torch.tensor of sinusoids with different frequencies.
+
+    Args:
+      seq_length: an optional Python int defining the output sequence length.
+        if the `position` argument is specified.
+      position:   [B, seq_length], optional position for each token in the
+        sequence, only required when the sequence is packed.
+
+    Returns:
+      [B, seqlen, D] if `position` is specified, else [1, seqlen, D]
+    """
+    if position is None:
+      assert seq_length is not None
+      # [1, seqlen]
+      position = torch.arange(seq_length, dtype=torch.float32)[None, :]
+    else:
+      assert position.ndim == 2, position.shape
+
+    num_timescales = self.embedding_dim // 2
+    log_timescale_increment = torch.log(
+        torch.tensor(float(self.max_timescale) / float(self.min_timescale))
+    ) / torch.maximum(
+        torch.tensor(num_timescales, dtype=torch.float32) - 1, torch.tensor(1)
+    )
+    inv_timescales = self.min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float32)
+        * -log_timescale_increment
+    )
+    scaled_time = position[:, :, None] * inv_timescales[None, None, :]
+    signal = torch.cat((torch.sin(scaled_time), torch.cos(scaled_time)), dim=2)
+    # Force usage of `np` rather than `jnp` to compute static values at trace
+    # time.
+    signal = F.pad(signal, (0, self.embedding_dim % 2, 0, 0, 0, 0))
+    return signal
+
+
+class MlpBlockWithMask(nn.Module):
+  """Transformer MLP / feed-forward block that supports masking."""
+
+  def __init__(
+      self,
+      mlp_dim: int,
+      d_model: int,
+      use_bias: bool = True,
+      dtype: torch.dtype = torch.float32,
+      activation_fn: nn.Module = nn.GELU,
+  ):
+    super().__init__()
+
+    self.mlp_dim = mlp_dim
+    self.d_model = d_model
+    self.use_bias = use_bias
+    self.dtype = dtype
+    self.activation_fn = activation_fn
+
+    self.c_fc = nn.Linear(
+        in_features=self.d_model,
+        out_features=self.mlp_dim,
+        dtype=self.dtype,
+        bias=self.use_bias,
+    )
+    self.c_proj = nn.Linear(
+        in_features=self.mlp_dim,
+        out_features=self.d_model,
+        dtype=self.dtype,
+        bias=self.use_bias,
+    )
+
+  def __call__(
+      self, inputs: torch.Tensor, mlp_mask: torch.Tensor
+  ) -> torch.Tensor:
+    """Applies Transformer MlpBlock with mask module."""
+    x = self.c_fc(inputs)
+    x = self.activation_fn()(x)
+    x = x * mlp_mask[..., None]  # First masking.
+    x = self.c_proj(x)
+    x = x * mlp_mask[..., None]  # Second masking.
+    return x
+
+
+class ResidualAttentionBlock(nn.Module):
+  """Transformer residual attention block."""
+
+  def __init__(
+      self,
+      d_model: int,
+      n_head: int,
+      mlp_dim: int,
+      dtype: torch.dtype = torch.float32,
+  ):
+    super().__init__()
+    self.d_model = d_model
+    self.n_head = n_head
+    self.mlp_dim = mlp_dim
+    self.dtype = dtype
+
+    self.attn = nn.MultiheadAttention(d_model, n_head, dtype=self.dtype)
+    self.ln_1 = nn.LayerNorm(d_model, dtype=self.dtype)
+    self.mlp = MlpBlockWithMask(
+        self.mlp_dim,
+        d_model,
+        use_bias=True,
+        dtype=self.dtype,
+        activation_fn=nn.ReLU,
+    )
+    self.ln_2 = nn.LayerNorm(d_model, dtype=self.dtype)
+
+  def attention(self, x: torch.Tensor, mask: torch.Tensor):
+    attn_mask = (
+        mask[:, None, None, :]
+        .repeat(1, self.n_head, x.shape[0], 1)
+        .flatten(0, 1)
+    )
+    attn_mask[attn_mask == 0] = float('-inf')
+    attn_mask[attn_mask == 1] = 0
+    return self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
+
+  def forward(self, x: torch.Tensor, mask: torch.Tensor):
+    x = x + self.attention(self.ln_1(x), mask.permute(1, 0))
+    x = x + self.mlp(self.ln_2(x), mask)
+    return x, mask
+
+
+class SequentialMultiInput(nn.Sequential):
+  """Sequential module that can take multiple inputs."""
+
+  def forward(self, *inputs):
+    for module in self._modules.values():
+      if isinstance(inputs, tuple):
+        inputs = module(*inputs)
+      else:
+        inputs = module(inputs)
+    return inputs
+
+
+class Transformer(nn.Module):
+  """Transformer implementation."""
+
+  def __init__(
+      self,
+      width: int,
+      layers: int,
+      heads: int,
+      mlp_dim: int,
+      dtype: torch.dtype = torch.float32,
+  ):
+    super().__init__()
+    self.width = width
+    self.layers = layers
+    self.heads = heads
+    self.mlp_dim = mlp_dim
+    self.dtype = dtype
+
+    self.resblocks = SequentialMultiInput(*[
+        ResidualAttentionBlock(self.width, self.heads, self.mlp_dim, self.dtype)
+        for _ in range(self.layers)
+    ])
+
+
+  def _concat_mask(self, mask: torch.Tensor, n_deep: int, after_bos: bool) -> torch.Tensor:
+      if mask is None:
+          return None
+      if mask.dim() == 2:                        # (B, L), 1/True = keep
+          B, L = mask.shape
+          deep_vis = torch.ones(B, n_deep, dtype=mask.dtype, device=mask.device)
+          if after_bos and L >= 1:
+              return torch.cat([mask[:, :1], deep_vis, mask[:, 1:]], dim=1)
+          else:
+              return torch.cat([deep_vis, mask], dim=1)
+      elif mask.dim() == 4:                      # (B, 1, 1, L), additive mask (0 keep)
+          B, _, _, L = mask.shape
+          deep_vis = torch.zeros(B, 1, 1, n_deep, dtype=mask.dtype, device=mask.device)
+          if after_bos and L >= 1:
+              return torch.cat([mask[..., :1], deep_vis, mask[..., 1:]], dim=-1)
+          else:
+              return torch.cat([deep_vis, mask], dim=-1)
+      else:
+          raise ValueError(f"Unsupported mask shape: {mask.shape}")
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      mask: torch.Tensor,
+      deep_tokens: nn.ParameterList = None,
+      *,
+      after_bos: bool = False  # place prompts after BOS/CLS at position 0
+  ) -> torch.Tensor:
+      """
+      If deep_tokens is provided, insert n_deep tokens at the input of layers
+      [insert_start_layer, insert_start_layer + len(deep_tokens)) and strip them after each block.
+      """
+      # fast path: no deep tokens
+      if not deep_tokens:
+          for block in self.resblocks:
+              x, mask = block(x, mask)
+          return x
+
+      # shapes/sizes
+      B, L, D = x.shape
+      n_deep = deep_tokens[0].shape[0]  # number of deep tokens per tuned layer
+      d_deep = len(deep_tokens)         # number of layers to tune
+      start = 1 # max(0, int(self.insert_start_layer))
+      end = min(start + d_deep, len(self.resblocks))
+
+      # sanity checks
+      for t in deep_tokens:
+          assert t.shape == (n_deep, D), f"Each deep token tensor must be (n_deep={n_deep}, D={D}), got {t.shape}"
+
+      for i, block in enumerate(self.resblocks):
+          if start <= i < end:
+              # select layer-specific deep tokens and expand across batch
+              pt = deep_tokens[i - start].unsqueeze(0).expand(B, -1, -1)  # (B, n_deep, D)
+
+              # place after BOS/CLS if present, otherwise at pure prefix
+              if after_bos and L >= 1:
+                  x_ext = torch.cat([x[:, :1], pt, x[:, 1:]], dim=1)      # [BOS] + deep + tokens
+              else:
+                  x_ext = torch.cat([pt, x], dim=1)                        # deep + tokens
+
+              mask_ext = self._concat_mask(mask, n_deep, after_bos=(after_bos and L >= 1))
+              y, _ = block(x_ext, mask_ext)
+              if after_bos and L >= 1:
+                  x = torch.cat([y[:, :1], y[:, 1 + n_deep:]], dim=1)
+              else:
+                  x = y[:, n_deep:, :]
+          else:
+              x, _ = block(x, mask)
+
+      return x
+
+def squeeze_multiple(tensor: torch.Tensor, dims: list) -> torch.Tensor:
+    """Dynamically squeezes specified singleton dimensions from a tensor.
+    
+    Args:
+        tensor (torch.Tensor): Input tensor.
+        dims (list): List of dimensions to squeeze.
+    
+    Returns:
+        torch.Tensor: Squeezed tensor.
+    """
+    for dim in sorted(dims, reverse=True):  # Sort in reverse to avoid shifting indices
+        if tensor.shape[dim] == 1:  # Only squeeze if dim is 1
+            tensor = tensor.squeeze(dim=dim)
+    return tensor
+  
+class GlobalAvgPooling(nn.Module):
+  """Performs a simple global pooling over the input with optional paddings.
+
+  Attributes:
+    pooling_dims: A list of dims to perform pooling over.
+    keepdims: If True, keep dimension of inputs after pooling.
+  """
+
+  pooling_dims: t.Sequence[int]
+  epsilon: float = 1e-8
+
+  def __init__(
+      self, pooling_dims: t.Sequence[int], epsilon: float = 1e-8
+  ):
+    super().__init__()
+    self.pooling_dims = pooling_dims
+    self.epsilon = epsilon
+
+    if not all([p_dims >= 0 for p_dims in self.pooling_dims]):
+      raise ValueError('pooling_dims must be non-negative integers.')
+
+  def __call__(
+      self,
+      inputs: torch.tensor,
+      compatible_paddings: torch.tensor,
+  ):
+    """Applies global average spatial pooling to inputs.
+
+    Args:
+      inputs: An input tensor.
+      compatible_paddings: paddings of inputs with shapes compatible with
+        inputs, e.g. compatible_paddings with shape [B, 1] for inputs with shape
+        [B, D].
+
+    Returns:
+      Output tensor with global pooling applied.
+    """
+    padded_value = torch.zeros_like(inputs)
+    padded_value = torch.ones_like(inputs) * padded_value
+    inputs = torch.where(compatible_paddings > 0, padded_value, inputs)
+    valid_inputs = (
+        torch.sum(
+            1.0 - compatible_paddings,
+            self.pooling_dims,
+            keepdims=True,
+            dtype=inputs.dtype,
+        )
+        + self.epsilon
+    )
+    inputs_sum = torch.sum(inputs, self.pooling_dims, keepdims=True)
+    outputs = torch.divide(inputs_sum, valid_inputs).type(inputs.dtype)
+    outputs = torch.squeeze(outputs, axis=self.pooling_dims)
+    # outputs = squeeze_multiple(outputs, self.pooling_dims)
+    return outputs
+
+
+class TextEncoder(nn.Module):
+  """Text encoder implementation."""
+
+  def __init__(
+      self,
+      config: t.Dict[str, int],
+      vocab_size: int,
+      dtype: torch.dtype = torch.float32,
+      scale_sqrt_depth: bool = True,
+  ):
+    super().__init__()
+    self.vocab_size = vocab_size
+    self.dtype = dtype
+    self.scale_sqrt_depth = scale_sqrt_depth
+
+    # The text tower layers are fixed independent of vision tower size.
+    self.transformer_layers = config['num_layers']
+    self.embedding_dim = config['hidden_size']
+    self.transformer_width = config['hidden_size']
+    self.mlp_dim = config['mlp_dim']
+    self.transformer_heads = config['num_heads']
+
+    self.token_embedding = nn.Embedding(
+        self.vocab_size, self.embedding_dim, dtype=self.dtype
+    )
+    self.pos_embedder = PositionalEmbedding(embedding_dim=self.embedding_dim)
+    self.transformer = Transformer(
+        width=self.transformer_width,
+        layers=self.transformer_layers,
+        heads=self.transformer_heads,
+        mlp_dim=self.mlp_dim,
+        dtype=self.dtype,
+    )
+    self.pooling = GlobalAvgPooling(pooling_dims=[1])
+    self.ln_final = nn.LayerNorm(self.transformer_width, dtype=self.dtype)
+
+  def __call__(
+      self,
+      ids: torch.tensor,
+      paddings: torch.tensor,
+      learnable_prompts: torch.Tensor = None,  # New parameter for learnable prompts
+      learning_method: str = None, # But addition can also be used [concat, sumate, entire_learnable, None]
+      deep_parameters: torch.nn.ParameterList = None,
+      device = 'cuda',
+  ):
+    # """Applies TextEncoder module."""
+    # _, seq_length = ids.shape
+
+    # x = self.token_embedding(ids)
+    # if self.scale_sqrt_depth:
+    #   x = x * (self.embedding_dim**0.5)
+
+    # mask = (paddings == 0).type(torch.float32)
+    # mask = mask.permute(1, 0)  # NL -> LN
+
+    # x = x + self.pos_embedder(seq_length=seq_length).to(device)
+    # x = x.permute(1, 0, 2)  # NLD -> LND
+    # x = self.transformer(x, mask)
+    # x = x.permute(1, 0, 2)  # LND -> NLD
+    # x = self.ln_final(x)
+    # x = self.pooling(x, compatible_paddings=paddings[:, :, None])
+    # return x
+
+    """Applies TextEncoder module with optional prompt learning (fixed input length 64)."""
+    batch_size, original_seq_length = ids.shape
+
+    x = self.token_embedding(ids)  # [B, L, D]
+
+    if self.scale_sqrt_depth:
+        x = x * (self.embedding_dim ** 0.5)
+
+    x = x.to(device)
+    if learnable_prompts is not None:
+      if learning_method == 'concat':
+        prompts = learnable_prompts.unsqueeze(0).expand(batch_size, -1, -1).to(x.device)  # [B, P, D]
+        x = torch.cat([prompts, x], dim=1)  # [B, P+L, D]
+        paddings = torch.cat([
+            torch.zeros((batch_size, prompts.size(1)), device=x.device), 
+            paddings
+        ], dim=1)
+
+      elif learning_method == 'sumate':
+        prompt_len = learnable_prompts.size(0)
+        x[:, :prompt_len, :] += learnable_prompts.unsqueeze(0)
+
+      elif learning_method == 'entire_learnable':
+        x = learnable_prompts.unsqueeze(0).expand(batch_size, -1, -1)
+        paddings = torch.zeros((batch_size, x.size(1)), device=x.device)
+
+    # 🔒 Ensure fixed sequence length (truncate or pad)
+    current_len = x.size(1)
+    if current_len > original_seq_length:
+      x = x[:, :original_seq_length, :]
+      paddings = paddings[:, :original_seq_length]
+    elif current_len < original_seq_length:
+      pad_len = original_seq_length - current_len
+      pad_embed = torch.zeros((batch_size, pad_len, x.size(2)), device=x.device)
+      pad_mask = torch.ones((batch_size, pad_len), device=x.device)  # masked-out padding
+      x = torch.cat([x, pad_embed], dim=1)
+      paddings = torch.cat([paddings, pad_mask], dim=1)
+
+    # Positional encoding and transformer
+    x = x + self.pos_embedder(seq_length=original_seq_length).to(device)
+    x = x.permute(1, 0, 2)  # [L, B, D]
+    mask = (paddings == 0).float().permute(1, 0)  # [L, B]
+    x = self.transformer(x, mask, deep_tokens=deep_parameters)
+    x = x.permute(1, 0, 2)  # [B, L, D]
+    x = self.ln_final(x)
+    x = self.pooling(x, compatible_paddings=paddings[:, :, None])
+
+    return x

Tipsomaly/reproduce.sh

@@ -0,0 +1,22 @@
+# This script provides a sample command to test given checkpoints on a dataset
+# For testing multiple datasets in a loop, refer to script train_test.sh
+
+models_dir="/kaggle/working/tips"
+data_root_dir="/kaggle/working/datasets"
+model_version='l14h' # model version name, like s14h, l14h, g14h for TIPS model and google/siglip2-large-patch16-512 for SigLIP2 model 
+checkpoint_path="./workspaces/trained_on_mvtec_default/vegan-arkansas/checkpoints"
+# visa checkpoint at: './workspaces/trained_on_visa_default/vegan-arkansas/checkpoints'
+
+### Test using industrial fixed prompts
+
+# test on VisA
+python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path --data_root_dir $data_root_dir --dataset visa --epoch 2 --model_version $model_version --fixed_prompt_type industrial
+
+# test on MVTec
+# python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path --data_root_dir $data_root_dir --dataset mvtec --epoch 2 --model_version $model_version --fixed_prompt_type industrial
+
+# test on a medical dataset with industrial prompts
+# python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path  --data_root_dir $data_root_dir --dataset cvc-colondb --fixed_prompt_type industrial --epoch 2 --model_version $model_version
+
+# test on a medical dataset with medical prompts
+# python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path  --data_root_dir $data_root_dir --dataset headct --fixed_prompt_type medical --epoch 2 --model_version $model_version

Tipsomaly/requirements.txt

@@ -0,0 +1,42 @@
+accelerate==1.11.0
+humanhash3==0.0.6
+humanize==4.14.0
+torch==2.8.0+cu126
+torchao==0.10.0
+torchaudio==2.8.0+cu126
+torchdata==0.11.0
+torchinfo==1.8.0
+torchmetrics==1.8.2
+torchsummary==1.5.1
+torchtune==0.6.1
+torchvision==0.23.0+cu126
+transformers==5.0.0    # NOTE: This version is required only for SigLIP2 experiments
+opencv-contrib-python==4.12.0.88
+opencv-python==4.12.0.88
+opencv-python-headless==4.12.0.88
+sentencepiece==0.2.1
+timm==1.0.20
+pillow==11.3.0
+scikit-image==0.25.2
+scikit-learn==1.6.1
+scikit-learn-intelex==2025.10.0
+scikit-multilearn==0.2.0
+scikit-optimize==0.10.2
+scikit-plot==0.3.7
+scikit-surprise==1.1.4
+scipy==1.15.3
+sklearn-pandas==2.2.0
+numpy==2.0.2
+pandas==2.2.2
+pandas-datareader==0.10.0
+pandas-gbq==0.29.2
+pandas-profiling==3.6.6
+pandas-stubs==2.2.2.240909
+pandasql==0.7.3
+tqdm==4.67.1
+PyYAML==6.0.3
+einops==0.8.1
+huggingface_hub==1.4.0
+
+# this is a curated and condensed list of the result of pip freeze > requirements.txt
+# most of our experiments was conducted on Kaggle on Python 3.11 and 3.12

Tipsomaly/test.py

@@ -0,0 +1,336 @@
+import hashlib
+import humanhash
+import os, sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "model"))
+
+import torch
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+from torchvision import transforms
+
+import random
+import numpy as np
+import pandas as pd 
+from pathlib import Path
+
+from tabulate import tabulate
+from tqdm import tqdm
+import subprocess
+import argparse
+from scipy.ndimage import gaussian_filter
+
+from model import tips, omaly
+from datasets import input_transforms, dataset
+from utils.metrics import image_level_metrics, pixel_level_metrics
+from utils.visualize import visualizer
+from utils.logger import get_logger, read_train_args
+from transformers import AutoProcessor, AutoModel, AutoTokenizer, SiglipTextModel, SiglipVisionModel
+from model.siglip2.siglip2_prompt_learnable import SiglipTextModelWithPromptLearning
+
+####################
+
+from model.big_vision import load_siglip
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def calc_soft_score(vis_feat, txt_feat, temp):
+    return F.softmax((vis_feat @ txt_feat.permute(0, 2, 1))/temp, dim=-1)
+
+def calc_sigm_score(vis_feat, txt_feat, temp, bias):
+    if vis_feat.dim() < 3:
+        vis_feat = vis_feat.unsqueeze(dim=1)
+    tempered_logits = vis_feat @ txt_feat.permute(0, 2, 1) * temp
+    probs = 1 / (1 + np.exp(-tempered_logits - bias))
+    return F.softmax(probs, dim=-1)
+
+def calc_sigm_score_hf(vis_feat, txt_feat, temp, bias):
+    if vis_feat.dim() < 3:
+        vis_feat = vis_feat.unsqueeze(dim=1)
+    logits = vis_feat @ txt_feat.permute(0, 2, 1) * temp + bias
+    probs = torch.sigmoid(logits)
+    return probs
+    
+def regrid_upsample_smooth(flat_scores, size, sigma):
+    h_w = int(flat_scores.shape[1] ** 0.5) 
+    regrided = flat_scores.reshape(flat_scores.shape[0], h_w, h_w, -1).permute(0, 3, 1, 2)
+    upsampled = torch.nn.functional.interpolate(regrided, (size, size), mode='bilinear').permute(0, 2, 3, 1)
+    rough_maps = (1-upsampled[..., 0] + upsampled[..., 1])/2
+    assert (rough_maps >= 0).all() and (rough_maps <= 1).all(), "All elements of rough_maps must be between 0 and 1"
+    anomaly_map = torch.stack([torch.from_numpy(gaussian_filter(map, sigma=sigma)) for map in rough_maps.detach().cpu()], dim=0)
+    return anomaly_map
+
+def create_tips(args, device):
+    # load dataset 
+    transform, target_transform = input_transforms.create_transforms_tips(args.image_size)
+
+    # load model
+    vision_encoder, text_encoder, tokenizer, temperature = tips.load_model.get_model(args.models_dir, args.model_version)
+    return vision_encoder.to(device), text_encoder.to(device), text_encoder.transformer.width, tokenizer, transform, target_transform, temperature
+
+# L/14, 512 
+def create_siglip2(args, device):
+    transform, target_transform = load_siglip.create_preprocessors_siglip2(args.image_size)
+    vision_encoder, text_encoder, tokenizer = load_siglip.build_siglip_modules(args.model_version, args.image_size)
+    # model.to(device)
+
+    temperature, bias = text_encoder.params['t'], text_encoder.params['b']
+    temperature = np.exp(torch.from_numpy(np.array(temperature)))
+    return vision_encoder, text_encoder, text_encoder.model.out_dim[1], tokenizer, transform, target_transform, temperature, bias
+
+def create_siglip2_hf(args, device):
+    tokenizer = AutoTokenizer.from_pretrained(args.model_version)
+    model = AutoModel.from_pretrained(args.model_version)
+    text_encoder = SiglipTextModelWithPromptLearning.from_pretrained(args.model_version).to(device)
+    vision_encoder = SiglipVisionModel.from_pretrained(args.model_version).to(device)
+    processor = AutoProcessor.from_pretrained(args.model_version)
+    def transform(x):
+        d = processor(images=x, return_tensors="pt")
+        return d['pixel_values'].squeeze(0)
+    target_transform = transforms.Compose([
+        transforms.Resize((args.image_size, args.image_size)),
+        transforms.ToTensor(),
+    ])
+    bias = model.logit_bias.to(device)
+    temperature = model.logit_scale.to(device).exp()
+    return vision_encoder, text_encoder, model.text_model.embeddings.token_embedding.embedding_dim, tokenizer, transform, target_transform, temperature, bias
+
+def test(args):
+    logger = get_logger(args.save_path)
+    # load dataset 
+
+    device = args.device
+    if args.backbone_name == 'tips':
+        bb_vision_encoder, bb_text_encoder, text_embd_dim, tokenizer, transform, target_transform, temperature = create_tips(args, device)
+        calc_score = lambda vis_feat, txt_feat: calc_soft_score(vis_feat, txt_feat, temperature)
+
+    elif args.backbone_name == 'siglip2':
+        bb_vision_encoder, bb_text_encoder, text_embd_dim, tokenizer, transform, target_transform, temperature, bias  = create_siglip2(args, device)
+        calc_score = lambda vis_feat, txt_feat: calc_sigm_score(vis_feat, txt_feat, temperature, bias)
+
+    elif args.backbone_name == 'siglip2-hf':
+        bb_vision_encoder, bb_text_encoder, text_embd_dim, tokenizer, transform, target_transform, temperature, bias = create_siglip2_hf(args, device)
+        calc_score = lambda vis_feat, txt_feat: calc_sigm_score_hf(vis_feat, txt_feat, temperature, bias)
+
+    text_encoder = omaly.text_encoder(tokenizer, bb_text_encoder, args.backbone_name, text_embd_dim, 64, args.prompt_learn_method, args.fixed_prompt_type, args.n_prompt, args.n_deep_tokens, args.d_deep_tokens)
+    vision_encoder = omaly.vision_encoder(bb_vision_encoder, args.backbone_name)
+
+    # class_names = desc.dataset_dict[args.dataset]
+    test_data = dataset.Dataset(args.data_path, transform, target_transform, args)
+    test_loader = DataLoader(test_data, batch_size=args.batch_size, num_workers=4, shuffle=False)
+    # test_loader = DataLoader(test_data, batch_size=8, shuffle=False, num_workers=1, prefetch_factor=2, pin_memory=True)
+    fixed_class_names = [clss.replace('_', ' ') for clss in test_data.cls_names]
+    
+    # extract features
+    with torch.no_grad():
+        # Fixed prototypes
+        fixed_text_features = text_encoder(fixed_class_names, device, learned=False)
+        fixed_text_features = fixed_text_features / fixed_text_features.norm(dim=-1, keepdim=True) # NOTE: For test also 
+    cls_text_features, seg_text_features = fixed_text_features, fixed_text_features
+    
+    if args.checkpoint_path:
+        assert not args.prompt_learn_method == 'none', 'The prompt_learn_method should not be none'
+        checkpoint = torch.load(args.params_path, weights_only=False)
+        text_encoder.learnable_prompts = checkpoint["learnable_prompts"] if isinstance(checkpoint, dict) else checkpoint
+
+        # text_encoder.learnable_prompts = chekpoint
+        # text_encoder.deep_parameters = chekpoint["deep_parameters"]
+        
+        learnable_class_names = ['object']
+        learnable_class_ids = torch.tensor([0])
+        print('The learnable prompts are read')
+        
+        # extract features
+        with torch.no_grad():
+            # Learnable prototypes
+            learnable_text_features = text_encoder(learnable_class_names, device, learned=True) # NOTE: important learned=True
+            learnable_text_features = learnable_text_features / learnable_text_features.norm(dim=-1, keepdim=True) # NOTE: For test also 
+        cls_text_features, seg_text_features = learnable_text_features, learnable_text_features
+    
+    if args.checkpoint_path and args.decoupled_prompt:
+        cls_text_features, seg_text_features = fixed_text_features, learnable_text_features
+        
+    dataset_preds = {cls_id: {'name': test_loader.dataset.cls_names[cls_id], 'img_scrs': [], 'img_lbls': [], 'pxl_scrs': [], 'pxl_lbls': [], 'paths': []} for cls_id in test_loader.dataset.class_ids}
+    for batch in tqdm(test_loader, desc="Extracting features", unit="batch"):
+        image = batch['img'].to(device)
+        label = batch['anomaly'].long().to(device)
+        abnorm_mask = batch['abnorm_mask'].squeeze(dim=1).to(device)
+        path = batch['img_path']
+        
+        # Indecies
+        cls_class_ids, seg_class_ids = batch['cls_id'], batch['cls_id']
+        if args.checkpoint_path and args.decoupled_prompt:
+            seg_class_ids = learnable_class_ids
+        elif args.checkpoint_path and not args.decoupled_prompt:
+            cls_class_ids, seg_class_ids = learnable_class_ids, learnable_class_ids
+
+        with torch.no_grad():
+            vision_features = vision_encoder(image)
+            vision_features = [feature / feature.norm(dim=-1, keepdim=True) for feature in vision_features] # NOTE: for test also
+
+            # calculate normal/abnormal scores
+            img_scr0 = calc_score(vision_features[0], cls_text_features[cls_class_ids]).squeeze(dim=1).detach() # prompt_class_ids cls_ids
+            img_scr1 = calc_score(vision_features[1], cls_text_features[cls_class_ids]).squeeze(dim=1).detach()
+
+            img_map = calc_score(vision_features[2], seg_text_features[seg_class_ids])
+            if args.aggregate_local2global:
+                max_local = torch.max(img_map, dim=1)[0]
+                img_scr0 = img_scr0 + max_local
+                img_scr1 = img_scr1 + max_local
+                
+            pxl_scr = regrid_upsample_smooth(img_map.detach(), args.image_size, args.sigma)
+            
+        for idx, cls_id in enumerate(batch['cls_id'].cpu().numpy()):
+            dataset_preds[cls_id]['img_scrs'].append([img_scr0[idx][1].cpu(), img_scr1[idx][1].cpu()])
+            dataset_preds[cls_id]['img_lbls'].append(label[idx].cpu())
+            dataset_preds[cls_id]['pxl_scrs'].append(pxl_scr[idx].cpu())
+            dataset_preds[cls_id]['pxl_lbls'].append(abnorm_mask[idx].cpu())
+            dataset_preds[cls_id]['paths'].append(path[idx])
+        
+    # calculate metrics
+    header = ['objects']+args.pixel_metrics+[mtr for mtr in args.image_metrics for _ in range(2)]
+    dataset_results = []
+    for cls_id in dataset_preds.keys():
+        cls_results = [dataset_preds[cls_id]['name']]
+        img_prds = np.array(dataset_preds[cls_id]['img_scrs'])
+        img_lbls = np.array(dataset_preds[cls_id]['img_lbls'])
+        pxl_prds = torch.stack(dataset_preds[cls_id]['pxl_scrs'], dim=0)
+        pxl_lbls = torch.stack(dataset_preds[cls_id]['pxl_lbls'], dim=0)
+        print(f'pxl_prds: ({pxl_prds.max()}, {pxl_prds.min()})')
+        print(f'img_prds: ({img_prds.max()}, {img_prds.min()})')
+
+        for px_mtr in args.pixel_metrics:
+            if not px_mtr == '':
+                cls_results.append(pixel_level_metrics(device, pxl_prds, pxl_lbls, px_mtr)*100)
+            
+        for im_mtr in args.image_metrics:
+            for col in range(img_prds.shape[1]):
+                cls_results.append(image_level_metrics(img_prds[:, col], img_lbls, im_mtr)*100)
+        
+        if args.visualize:
+            img_path = f"{args.dataset}/{dataset_preds[cls_id]['name']}"
+            visualizer(dataset_preds[cls_id]['paths'], pxl_prds.cpu().numpy(), pxl_lbls.cpu().numpy(), args.image_size, img_path, save_path=f'{args.save_path}/img/', draw_contours=True)
+            
+        dataset_results.append(cls_results)
+        
+    df = pd.DataFrame(dataset_results, columns=header)
+    mean_values = ['Mean'] + df.iloc[:, 1:].mean().tolist()
+    # df = df.append(mean_values)
+    df.loc[len(df)] = mean_values
+    df = df.round(2)
+    
+    # store the results
+    results_text = tabulate(df, headers='keys', tablefmt='pretty') 
+    logger.info(results_text)
+    
+def make_human_readable_name(args, exclude=['model_name', 'dataset', 'dataset_category', 'data_path',
+                                            'checkpoint_path', 'training_path', "Timestamp",
+                                            "metrics", "device", "available_devices", "epochs", "visualize", 'help', None]):
+    args=vars(args)
+    name_value_pairs = [
+        f"{k}_{v}"
+        for k,v in args.items()
+        if k not in exclude # Exclude "help" or invalid arguments
+    ]   
+    combined = ",".join(sorted(name_value_pairs))  # Sorting ensures consistent order
+    hash_value = hashlib.sha256(combined.encode()).hexdigest()
+    human_hash = humanhash.humanize(hash_value, words=2) # 'hash' #
+    return human_hash
+
+def str2bool(v):
+    if isinstance(v, bool):
+       return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser("TIPSomaly", add_help=True)
+    # model
+    parser.add_argument("--image_size", type=int, default=518, help="image size") #224 if is_low_res else 448
+    parser.add_argument("--seed", type=int, default=111, help="random seed")
+
+    parser.add_argument("--metrics", type=str, default='image-pixel-level')
+    parser.add_argument("--device", type=str, default="cuda", help="type of device, can be cuda or cpu")
+    parser.add_argument("--available_devices", type=int, nargs='+', default=[0, 1, 2, 3, 4, 5, 6, 7], help="array of possible cuda devices")
+    parser.add_argument("--model_name", type=str, default="trained_on_visa_siglip2_both_mvtec", help="")
+    parser.add_argument("--models_dir", type=str, default="./tips", help="directory of the base model of tips")
+    parser.add_argument("--data_root_dir", type=str, default="./datasets", help="root directory for all datasets to be placed in")
+    parser.add_argument("--checkpoint_path", type=str, default='None', help="")
+    parser.add_argument("--epoch", type=int, default=1, help="")
+    parser.add_argument("--batch_size", type=int, default=8)
+
+    parser.add_argument("--sigma", type=int, default=4, help="zero shot")
+    
+    parser.add_argument("--dataset", type=str, default="visa")
+    parser.add_argument("--dataset_category", type=str, default='', help="train dataset categories")
+    
+    parser.add_argument("--class_name", type=str, nargs='+', default=['all'], help="train class name")
+    
+    parser.add_argument("--image_metrics", type=str, nargs='+', default=['auroc', 'ap', 'f1-max'], help="")
+    parser.add_argument("--pixel_metrics", type=str, nargs='+', default=['auroc', 'aupro', 'f1-max'], help="")
+
+    parser.add_argument("--k_shot", type=int, default=0, help="number of samples per class for few-shot learning. 0 means use all data.")
+    
+    parser.add_argument("--type", type=str, default='test') 
+    parser.add_argument("--visualize", type=str2bool, default=False)
+    parser.add_argument("--log_dir", type=str, default="")
+    
+    ##########################
+    ### Method Arguements ####
+    parser.add_argument("--backbone_name", type=str, default='tips', choices=["tips", "siglip2", "siglip2-hf"])
+    parser.add_argument("--model_version", type=str, default='l14h', choices=["s14h","b14h","l14h","so4h","g14l","g14h", \
+                                                                                "B/16", "L/16", "So400m/14", "So400m/16", "g-opt/16", \
+                                                                                "google/siglip2-so400m-patch16-256", "google/siglip2-large-patch16-512"])
+    
+    parser.add_argument("--n_deep_tokens", type=int, default=0)
+    parser.add_argument("--d_deep_tokens", type=int, default=0)
+    parser.add_argument("--n_prompt", type=int, default=8)
+    parser.add_argument("--fixed_prompt_type", type=str, default='industrial', choices=['industrial', 'medical', 'object_agnostic'])
+    
+    parser.add_argument("--prompt_learn_method", type=str, default='concat', choices=['concat', 'sumate', 'entire_learnable', 'none'])
+    parser.add_argument("--decoupled_prompt", type=str2bool, default=True)
+    parser.add_argument("--aggregate_local2global", type=str2bool, default=True)
+    
+    args = parser.parse_args()
+
+    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
+        os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(map(str, args.available_devices)) if len(args.available_devices) > 1 else str(args.available_devices[0])
+        command = [sys.executable, __file__, ] + sys.argv[1:] 
+        process = subprocess.Popen(command, env=os.environ)
+        process.wait()
+        
+    else:
+        setup_seed(args.seed)
+
+        #### ONLY KAGGLE
+        # args.dataset = f'{args.dataset}-ad'
+        # base_paths = [Path(p) for p in [f'{DATA_ROOT_DIR}/{args.dataset_category}/{args.dataset}/']]
+        # args.data_path = [str(next(p.iterdir())) for p in base_paths]
+        
+        args.data_path = [f'{args.data_root_dir}/{args.dataset_category}/{args.dataset}/']
+        if not args.checkpoint_path:
+            args.log_dir = make_human_readable_name(args)
+            args.save_path = f'./workspaces/{args.model_name}/{args.log_dir}/quantative/NoTrain/{args.dataset}'
+        else: # ./workspaces/test/blah-blah/checkpoints/
+            splits = args.checkpoint_path.split('/')            
+            args.params_path = f'{args.checkpoint_path}/learnable_params_{args.epoch}.pth'
+            args.model_name = splits[-3]
+            args.log_dir = splits[-2]
+            args.save_path = f'{"/".join(splits[:-1])}/quantative/epoch_{args.epoch}/{args.dataset}'
+            ### 
+            train_args = read_train_args(args.checkpoint_path) # ./workspaces/{args.model_name}/{args.log_dir}/args.txt
+            args.prompt_learn_method = train_args['prompt_learn_method']
+            assert not train_args['prompt_learn_method'] is None, 'prompt_learn_method should not be none'
+            
+        print(args)
+        print(f"Data Path: {args.data_path}, Log Directory: {args.log_dir}, Save Path: {args.save_path}")
+        test(args) 

Tipsomaly/train.py

@@ -0,0 +1,353 @@
+import os, sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "model"))
+import subprocess
+import argparse
+import hashlib
+import humanhash
+from collections import defaultdict
+from tqdm import tqdm
+
+import random
+import numpy as np
+from scipy.ndimage import gaussian_filter
+
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from torch.utils.data import DataLoader
+from transformers import AutoProcessor, AutoModel, AutoTokenizer, SiglipTextModel, SiglipVisionModel
+
+from datasets import input_transforms, dataset
+from utils.loss import FocalLoss, BinaryDiceLoss
+from utils.logger import save_args_to_file, get_logger
+from torch.utils.tensorboard import SummaryWriter
+
+from model import tips
+from model import omaly
+from model.big_vision import load_siglip
+from model.siglip2.siglip2_prompt_learnable import SiglipTextModelWithPromptLearning
+
+loss_names = {'img_ls_ce': 'LS CE', 'pxl_ls_fc': 'LS FC', \
+                'plx_ls_dc_p': 'LS DC P', 'plx_ls_dc_n': 'LS DC N', \
+                'emb_l1_nrm': 'LS L1 NRM', 'epc_ls': 'total'}
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def seed_worker(worker_id):
+    worker_seed = 111 + worker_id
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+def calc_soft_score(vis_feat, txt_feat, temp):
+    return F.softmax((vis_feat @ txt_feat.permute(0, 2, 1))/temp, dim=-1)
+
+def calc_sigm_score(vis_feat, txt_feat, temp, bias):
+    if vis_feat.dim() < 3:
+        vis_feat = vis_feat.unsqueeze(dim=1)
+    tempered_logits = vis_feat @ txt_feat.permute(0, 2, 1) * temp
+    probs = 1 / (1 + np.exp(-tempered_logits - bias))
+    return F.softmax(probs, dim=-1)
+
+def calc_sigm_score_hf(vis_feat, txt_feat, temp_non_exp, bias):
+    if vis_feat.dim() < 3:
+        vis_feat = vis_feat.unsqueeze(dim=1)
+    logits = vis_feat @ txt_feat.permute(0, 2, 1) * temp_non_exp.exp() + bias
+    probs = torch.sigmoid(logits)
+    return probs
+
+def create_tips(args, device):
+    # load dataset 
+    transform, target_transform = input_transforms.create_transforms_tips(args.image_size)
+
+    # load model
+    vision_encoder, text_encoder, tokenizer, temperature = tips.load_model.get_model(args.models_dir, args.model_version)
+    return vision_encoder.to(device), text_encoder.to(device), text_encoder.transformer.width, tokenizer, transform, target_transform, temperature
+
+def create_siglip2(args, device):
+    transform, target_transform = load_siglip.create_preprocessors_siglip2(args.image_size)
+    vision_encoder, text_encoder, tokenizer = load_siglip.build_siglip_modules(args.model_version, args.image_size)
+    # model.to(device)
+
+    temperature, bias = text_encoder.params['t'], text_encoder.params['b']
+    temperature = np.exp(torch.from_numpy(np.array(temperature)))
+    return vision_encoder, text_encoder, text_encoder.model.out_dim[1], tokenizer, transform, target_transform, temperature, bias
+
+def create_siglip2_hf(args, device):
+    tokenizer = AutoTokenizer.from_pretrained(args.model_version)
+    model = AutoModel.from_pretrained(args.model_version)
+    text_encoder = SiglipTextModelWithPromptLearning.from_pretrained(args.model_version).to(device)
+    vision_encoder = SiglipVisionModel.from_pretrained(args.model_version).to(device)
+    processor = AutoProcessor.from_pretrained(args.model_version)
+    def transform(x):
+        d = processor(images=x, return_tensors="pt")
+        return d['pixel_values'].squeeze(0)
+    target_transform = transforms.Compose([
+        transforms.Resize((args.image_size, args.image_size)),
+        transforms.ToTensor(),
+    ])
+    bias = model.logit_bias.to(device)
+    temp_non_exp = model.logit_scale.to(device)
+    return vision_encoder, text_encoder, model.text_model.embeddings.token_embedding.embedding_dim, tokenizer, transform, target_transform, temp_non_exp, bias
+
+def regrid_upsample_smooth(flat_scores, size, sigma):
+    upsampled = regrid_upsample(flat_scores, size)
+    anomaly_map = torch.stack([torch.from_numpy(gaussian_filter(map, sigma=sigma)) for map in upsampled.detach().cpu()], dim=0)
+    return anomaly_map
+
+def regrid_upsample(flat_scores, size):
+    h_w = int(flat_scores.shape[1] ** 0.5) 
+    regrided = flat_scores.reshape(flat_scores.shape[0], h_w, h_w, -1).permute(0, 3, 1, 2)
+    upsampled = torch.nn.functional.interpolate(regrided, (size, size), mode='bilinear').permute(0, 2, 3, 1)
+    return upsampled
+
+def turn_gradient_off(model):
+    print("Turning off gradients in both the image and the text encoder")
+    for _, param in model.named_parameters():
+        param.requires_grad_(False)
+
+    enabled = set()
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            enabled.add(name)
+    # print(f"Parameters to be updated: {enabled}")
+
+    model.eval()
+    return model
+
+def train(args):
+    epochs = args.epoch
+    device = args.device
+
+    writer = SummaryWriter(log_dir=args.experiment_root)
+    logger = get_logger(args.experiment_root)
+    
+    if args.backbone_name == 'tips':
+        bb_vision_encoder, bb_text_encoder, text_embd_dim, tokenizer, transform, target_transform, temperature = create_tips(args, device)
+        calc_score = lambda vis_feat, txt_feat: calc_soft_score(vis_feat, txt_feat, temperature)
+    elif args.backbone_name == "siglip2":
+        bb_vision_encoder, bb_text_encoder, text_embd_dim, tokenizer, transform, target_transform, temperature, bias  = create_siglip2(args, device)
+        calc_score = lambda vis_feat, txt_feat: calc_sigm_score(vis_feat, txt_feat, temperature, bias)
+    elif args.backbone_name == 'siglip2-hf':
+        bb_vision_encoder, bb_text_encoder, text_embd_dim, tokenizer, transform, target_transform, temperature, bias = create_siglip2_hf(args, device)
+        calc_score = lambda vis_feat, txt_feat: calc_sigm_score_hf(vis_feat, txt_feat, temperature, bias)
+
+    bb_text_encoder = bb_text_encoder.to(device)
+    bb_vision_encoder = bb_vision_encoder.to(device)
+    bb_text_encoder = turn_gradient_off(bb_text_encoder)
+    bb_vision_encoder = turn_gradient_off(bb_vision_encoder)
+    text_encoder = omaly.text_encoder(tokenizer, bb_text_encoder, args.backbone_name, text_embd_dim, 64, args.prompt_learn_method, args.fixed_prompt_type, args.n_prompt, args.n_deep_tokens, args.d_deep_tokens)
+    vision_encoder = omaly.vision_encoder(bb_vision_encoder, args.backbone_name)
+
+
+    # load dataset 
+    # class_names = desc.dataset_dict[args.dataset]
+    train_data = dataset.Dataset(args.data_path, transform, target_transform, args)    
+
+    g = torch.Generator()
+    g.manual_seed(args.seed)
+    train_loader = DataLoader(train_data, batch_size=args.batch_size, num_workers=4, shuffle=False)
+
+    # class_names = [clss.replace('_', ' ') for clss in train_data.cls_names]
+    # class_ids = train_data.class_ids
+    class_names = ['object']
+    class_ids = torch.tensor([0])
+    
+    # Define losses 
+    bce_loss = torch.nn.CrossEntropyLoss()
+    loss_focal = FocalLoss() 
+    loss_dice = BinaryDiceLoss()
+    
+    # Define optimizer
+    optimizer = torch.optim.Adam(
+        list(text_encoder.learnable_prompts),# + list(text_encoder.deep_parameters),
+        lr=args.learning_rate,
+        betas=(0.5, 0.999)
+    )
+    train_stats = defaultdict(list)
+
+    torch.autograd.set_detect_anomaly(True)
+    train_loader_cpu = [bat for bat in train_loader]
+    
+    global_step = 0
+
+    text_encoder.train() 
+    text_encoder.to(device)
+    vision_encoder.train() 
+    vision_encoder.to(device)
+    for epoch in range(epochs):  # Add epoch loop
+        print(f"Epoch {epoch + 1}/{epochs}")
+        epoch_loss = defaultdict(int)
+        
+        for batch in tqdm(train_loader_cpu, desc="Train", unit="batch"):
+            image = batch['img'].to(device)
+            # cls_ids = batch['cls_id']
+            label = batch['anomaly'].long().to(device)
+            abnorm_mask = batch['abnorm_mask'].squeeze(dim=1).to(device)
+            
+            # extract features
+            text_features = text_encoder(class_names, device, learned=True)
+            text_features = text_features / text_features.norm(dim=-1, keepdim=True) # NOTE: For test also 
+            with torch.no_grad():
+                vision_features = vision_encoder(image)
+                vision_features = [feature / feature.norm(dim=-1, keepdim=True) for feature in vision_features] # NOTE: for test also
+                
+            # calculate normal/abnormal scores (since TIPS has two global visual embeddings we have two calculated image-level scores)
+            img_scr0 = calc_score(vision_features[0], text_features[class_ids]).squeeze(dim=1)
+            img_scr1 = calc_score(vision_features[1], text_features[class_ids]).squeeze(dim=1)
+            
+            img_map = calc_score(vision_features[2], text_features[class_ids])
+            anomaly_map = regrid_upsample(img_map, args.image_size)
+            abnorm_mask[abnorm_mask > 0.5], abnorm_mask[abnorm_mask< 0.5] = 1, 0
+
+            # Calculate loss
+            anomaly_map = anomaly_map.permute(0, 3, 1, 2)
+            ls_fc = loss_focal(anomaly_map, abnorm_mask)
+            ls_dc_p = loss_dice(anomaly_map[:, 1, :, :], abnorm_mask)
+            ls_dc_n = loss_dice(anomaly_map[:, 0, :, :], 1-abnorm_mask)
+            
+            ls_cls = bce_loss(img_scr0, label) + bce_loss(img_scr1, label) 
+            ls_seg = ls_fc + ls_dc_p + ls_dc_n  # (pixel loss)
+            if args.cls_seg_los == 'both':      # (image loss)
+                loss_total = ls_cls + ls_seg
+            elif args.cls_seg_los == 'seg':
+                loss_total = ls_seg
+            elif args.cls_seg_los == 'cls':
+                loss_total = ls_cls
+
+            # L1 Regularization term
+            l1_norm = torch.sum(torch.abs(text_features))
+            loss_total = loss_total + l1_norm * args.l1_lambda
+
+            # Train
+            optimizer.zero_grad() 
+            loss_total.backward() 
+            optimizer.step()
+            
+            # log
+            epoch_loss['img_ls_ce'] += ls_cls.item()
+            epoch_loss['pxl_ls_fc'] += ls_fc.item()
+            epoch_loss['plx_ls_dc_p'] += ls_dc_p.item()
+            epoch_loss['plx_ls_dc_n'] += ls_dc_n.item()
+            epoch_loss['epc_ls'] += loss_total.item()
+            epoch_loss['emb_l1_nrm'] += l1_norm.item()
+
+            # Tensorboard update for each batch
+            writer.add_scalar(f"Loss/img_ls_ce", ls_cls.item(), global_step)
+            writer.add_scalar(f"Loss/pxl_ls_fc", ls_fc.item(), global_step)
+            writer.add_scalar(f"Loss/plx_ls_dc_p", ls_dc_p.item(), global_step)
+            writer.add_scalar(f"Loss/plx_ls_dc_n", ls_dc_n.item(), global_step)
+            writer.add_scalar(f"Loss/epc_ls", loss_total.item(), global_step)
+            writer.add_scalar(f"Loss/emb_l1_nrm", l1_norm.item(), global_step)
+            global_step += 1
+        
+        # Calc epoch mean loss
+        num_batches = len(train_loader)
+        for key, val in epoch_loss.items():
+            train_stats[key].append(val / num_batches)
+        
+        # Print mean losses at the end of the epoch
+        epoch_details = f"Epoch {epoch + 1} Mean Losses: "
+        for key, val in epoch_loss.items():
+            epoch_details = epoch_details + f"{loss_names[key]}: {train_stats[key][-1]:.4f}, " 
+        logger.info(epoch_details[:-2])
+        
+        torch.save({"learnable_prompts":text_encoder.learnable_prompts}, 
+                   f'{args.save_path}/learnable_params_{epoch+1}.pth')
+                    # "deep_parameters":text_encoder.deep_parameters}, 
+        print(f'checkpoints saved for epoch {epoch+1}.')
+        
+def make_human_readable_name(args, exclude=['model_name', 'dataset', 'dataset_category', 'epoch', 'data_path',
+                                            'checkpoint_path', 'training_path', "Timestamp",
+                                            "metrics", "device", "available_devices", "epochs", "visualize", 'help', None]):
+    args=vars(args)
+    name_value_pairs = [
+        f"{k}_{v}"
+        for k,v in args.items()
+        if k not in exclude # Exclude "help" or invalid arguments
+    ]   
+    combined = ",".join(sorted(name_value_pairs))  # Sorting ensures consistent order
+    hash_value = hashlib.sha256(combined.encode()).hexdigest()
+    human_hash = humanhash.humanize(hash_value, words=2)
+    return human_hash.replace('-', '_')
+
+def str2bool(v):
+    if isinstance(v, bool):
+       return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+if __name__ == '__main__':  
+
+    dss = ['mvtec']
+
+    parser = argparse.ArgumentParser("TIPSomaly", add_help=True)
+    # model
+    parser.add_argument("--image_size", type=int, default=518, help="image size")
+    parser.add_argument("--seed", type=int, default=111, help="random seed")
+
+    parser.add_argument("--epoch", type=int, default=5, help="epochs")
+    parser.add_argument("--learning_rate", type=float, default=0.001)
+
+    parser.add_argument("--metrics", type=str, default='image-pixel-level')
+    parser.add_argument("--device", type=str, default="cuda", help="type of device, can be cuda or cpu")
+    parser.add_argument("--available_devices", type=int, nargs='+', default=[0, 1, 2, 3, 4, 5, 6, 7], help="array of possible cuda devices")
+    parser.add_argument("--model_name", type=str, default="tips_test", help="cuda device")
+    parser.add_argument("--models_dir", type=str, default="./tips", help="directory of the base model of tips")
+    parser.add_argument("--data_root_dir", type=str, default="./datasets", help="root directory for all datasets to be placed in")
+    parser.add_argument("--batch_size", type=int, default=8)
+
+    parser.add_argument("--sigma", type=int, default=4, help="zero shot")
+    
+    parser.add_argument("--dataset", type=str, default="visa")
+    parser.add_argument("--dataset_category", type=str, default='', help="train dataset categories")
+    
+    parser.add_argument("--type", type=str, default='train') 
+    parser.add_argument("--class_name", type=str, nargs='+', default=['all'], help="train class name")
+    parser.add_argument("--k_shot", type=int, default=0, help="number of samples per class for few-shot learning. 0 means use all data.")
+
+    ##########################
+    ### Method Arguements ####
+    parser.add_argument("--model_version", type=str, default='l14h', choices=["s14h","b14h","l14h","so4h","g14l","g14h", \
+                                                                                "B/16", "L/16", "So400m/14", "So400m/16", "g-opt/16", \
+                                                                                    "google/siglip2-so400m-patch16-256", "google/siglip2-large-patch16-512"])
+    parser.add_argument("--n_deep_tokens", type=int, default=0)
+    parser.add_argument("--d_deep_tokens", type=int, default=0)
+    parser.add_argument("--n_prompt", type=int, default=8)
+    parser.add_argument("--fixed_prompt_type", type=str, default='industrial')
+    
+    parser.add_argument("--prompt_learn_method", type=str, default='concat', choices=['concat', 'sumate', 'entire_learnable', 'none'])
+    parser.add_argument("--cls_seg_los", type=str, default='seg', choices=['both', 'seg', 'cls'])
+    parser.add_argument("--l1_lambda", type=float, default=0.0)
+    parser.add_argument("--backbone_name", type=str, default='tips', choices=["tips", "siglip2", "siglip2-hf"])
+
+    args = parser.parse_args()        
+
+    command = [sys.executable, __file__, ] + sys.argv[1:] 
+    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
+        os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(map(str, args.available_devices)) if len(args.available_devices) > 1 else str(args.available_devices[0])
+        process = subprocess.Popen(command, env=os.environ)
+        process.wait()
+        
+    else:
+        print(args)
+        setup_seed(args.seed)
+        args.log_dir = make_human_readable_name(args)
+        args.data_path = [f'{args.data_root_dir}/{args.dataset_category}/{args.dataset}/']
+        args.experiment_root = f'./workspaces/trained_on_{args.dataset}_{args.model_name}/{args.log_dir}'
+        args.save_path = f'{args.experiment_root}/checkpoints'
+        os.makedirs(args.save_path, exist_ok=True)
+        
+        save_args_to_file(args, command) # ./workspaces/{args.model_name}/{args.log_dir}/args.txt
+        
+        print(f"Data Path: {args.data_path}, Log Directory: {args.log_dir}, Save Path: {args.save_path}")
+        train(args)

Tipsomaly/train_test.sh

@@ -0,0 +1,34 @@
+# This script provides the command to train your model
+
+# By setting "model_name" and "log_dir" you define the path for the checkpoints to be saved and 
+# you can use the same values later to test on other datasets in a loop easily  
+model_version='l14h' # l14h or l14
+model_name=$1
+epoch=$2
+
+models_dir="/kaggle/working/tips"
+data_root_dir="/kaggle/working/datasets"
+# checkpoint_path="/kaggle/working/checkpoints" # uncomment for testing results
+
+# Train on MVTec
+python train.py --models_dir $models_dir --model_name $model_name --data_root_dir $data_root_dir --dataset mvtec --cls_seg_los seg --l1_lambda 0.0 --d_deep_tokens 0 --n_deep_tokens 0 --epoch $epoch --model_version $model_version --fixed_prompt_type industrial
+
+# Train on VisA
+# python train.py --models_dir $models_dir --model_name $model_name --data_root_dir $data_root_dir --dataset visa --cls_seg_los seg --l1_lambda 0.0 --d_deep_tokens 0 --n_deep_tokens 0 --epoch $epoch --model_version $model_version --fixed_prompt_type industrial
+
+
+# Test multiple datasets in a loop
+
+# for dataset in visa mpdd btad sdd dagm dtd; do
+#     python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path --data_root_dir $data_root_dir --dataset $dataset --epoch $epoch --model_version $model_version --fixed_prompt_type industrial
+# done
+
+# Medical segmentation datasets - using learned prompts
+# for dataset in isic tn3k cvc-colondb cvc-clinicdb; do
+#     python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path --data_root_dir $data_root_dir --dataset $dataset --fixed_prompt_type industrial --epoch $epoch --model_version $model_version
+# done
+
+# Medical classification datasets - using medical prompts
+# for dataset in headct brainmri br35h; do
+#     python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path --data_root_dir $data_root_dir --dataset $dataset --fixed_prompt_type medical --epoch $epoch --model_version $model_version
+# done

Tipsomaly/utils/logger.py

@@ -0,0 +1,72 @@
+
+import logging
+import os
+import datetime
+
+def read_train_args(training_path):
+    # Check if args.txt exists in the training_path
+    args_file_path = os.path.join(training_path, 'args.txt')
+    configurations_dict = {}  # Dictionary to store configurations
+    last_config = {}
+    if os.path.exists(args_file_path):
+        with open(args_file_path, 'r') as f:
+            # Read the entire content of the file
+            file_content = f.read().strip()
+        
+        # Split the content into different configurations based on the 'Timestamp' keyword
+        configurations = file_content.split('Timestamp:')
+        
+        # Iterate over each configuration to populate the dictionary
+        for config in configurations:
+            if config.strip():
+                # Convert the configuration to a dictionary
+                file_args_dict = {}
+                for line in config.strip().split('\n'):
+                    if ':' in line:
+                        key, value = line.split(':', 1)
+                        file_args_dict[key.strip()] = value.strip()
+                
+                # Store the configuration in the dictionary with a unique key
+                timestamp = file_args_dict.get('Timestamp', 'Unknown')
+                configurations_dict[timestamp] = file_args_dict
+                
+        if configurations_dict:
+            last_timestamp = max(configurations_dict.keys())
+            last_config = configurations_dict[last_timestamp]
+    else:
+        print('train args does not exists')
+    return last_config
+
+def save_args_to_file(args, command, log_dir=''):
+    args_file_path = os.path.join(args.save_path, log_dir, 'args.txt')
+    os.makedirs(os.path.dirname(args_file_path), exist_ok=True)
+    if os.path.exists(args_file_path):
+        print(f"Warning: The file {args_file_path} already exists and will be overwritten.")
+    with open(args_file_path, 'a') as f:  # Change 'w' to 'a' to append to the file
+        f.write("\n")  # Add new line before writing to the file
+        f.write(f"Timestamp: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")  # Add timestamp
+        for arg, value in vars(args).items():
+            f.write(f"{arg}: {value}\n")
+        f.write(f"Command arguments: {' '.join(command)}\n")  # Add the command arguments to the file
+
+def get_logger(save_path):
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+
+    txt_path = os.path.join(save_path, 'log.txt')
+    # logger
+    root_logger = logging.getLogger()
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+    root_logger.setLevel(logging.WARNING)
+    logger = logging.getLogger('test')
+    formatter = logging.Formatter('%(asctime)s.%(msecs)03d - %(levelname)s: %(message)s',
+                                    datefmt='%y-%m-%d %H:%M:%S')
+    logger.setLevel(logging.INFO)
+    file_handler = logging.FileHandler(txt_path, mode='a')
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+    return logger

Tipsomaly/utils/loss.py

@@ -0,0 +1,107 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from math import exp
+
+class FocalLoss(nn.Module):
+    """
+    copy from: https://github.com/Hsuxu/Loss_ToolBox-PyTorch/blob/master/FocalLoss/FocalLoss.py
+    This is a implementation of Focal Loss with smooth label cross entropy supported which is proposed in
+    'Focal Loss for Dense Object Detection. (https://arxiv.org/abs/1708.02002)'
+        Focal_Loss= -1*alpha*(1-pt)*log(pt)
+    :param alpha: (tensor) 3D or 4D the scalar factor for this criterion
+    :param gamma: (float,double) gamma > 0 reduces the relative loss for well-classified examples (p>0.5) putting more
+                    focus on hard misclassified example
+    :param smooth: (float,double) smooth value when cross entropy
+    :param balance_index: (int) balance class index, should be specific when alpha is float
+    :param size_average: (bool, optional) By default, the losses are averaged over each loss element in the batch.
+    """
+
+    def __init__(self, apply_nonlin=None, alpha= None, gamma=2, balance_index=1, smooth=1e-5, size_average=True):
+        super(FocalLoss, self).__init__()
+        self.apply_nonlin = apply_nonlin
+        self.alpha = alpha
+        self.gamma = gamma
+        self.balance_index = balance_index
+        self.smooth = smooth
+        self.size_average = size_average
+
+        if self.smooth is not None:
+            if self.smooth < 0 or self.smooth > 1.0:
+                raise ValueError('smooth value should be in [0,1]')
+
+    def forward(self, logit, target):
+        if self.apply_nonlin is not None:
+            logit = self.apply_nonlin(logit)
+        num_class = logit.shape[1]
+
+        if logit.dim() > 2:
+            # N,C,d1,d2 -> N,C,m (m=d1*d2*...)
+            logit = logit.view(logit.size(0), logit.size(1), -1)
+            logit = logit.permute(0, 2, 1).contiguous()
+            logit = logit.view(-1, logit.size(-1))
+        target = torch.squeeze(target, 1)
+        target = target.view(-1, 1)
+        alpha = self.alpha
+
+        if alpha is None:
+            alpha = torch.ones(num_class, 1)
+        elif isinstance(alpha, (list, np.ndarray)):
+            assert len(alpha) == num_class
+            alpha = torch.FloatTensor(alpha).view(num_class, 1)
+            alpha = alpha / alpha.sum()
+        elif isinstance(alpha, float):
+            alpha = torch.ones(num_class, 1)
+            alpha = alpha * (1 - self.alpha)
+            alpha[self.balance_index] = self.alpha
+
+        else:
+            raise TypeError('Not support alpha type')
+
+        if alpha.device != logit.device:
+            alpha = alpha.to(logit.device)
+
+        idx = target.cpu().long()
+
+        one_hot_key = torch.FloatTensor(target.size(0), num_class).zero_()
+        one_hot_key = one_hot_key.scatter_(1, idx, 1)
+        if one_hot_key.device != logit.device:
+            one_hot_key = one_hot_key.to(logit.device)
+
+        if self.smooth:
+            one_hot_key = torch.clamp(
+                one_hot_key, self.smooth / (num_class - 1), 1.0 - self.smooth)
+        pt = (one_hot_key * logit).sum(1) + self.smooth
+        logpt = pt.log()
+
+        gamma = self.gamma
+
+        alpha = alpha[idx]
+        alpha = torch.squeeze(alpha)
+        loss = -1 * alpha * torch.pow((1 - pt), gamma) * logpt
+
+        if self.size_average:
+            loss = loss.mean()
+        return loss
+
+
+class BinaryDiceLoss(nn.Module):
+    def __init__(self):
+        super(BinaryDiceLoss, self).__init__()
+
+    def forward(self, input, targets):
+        # 获取每个批次的大小 N
+        N = targets.size()[0]
+        # 平滑变量
+        smooth = 1
+        # 将宽高 reshape 到同一纬度
+        input_flat = input.view(N, -1)
+        targets_flat = targets.view(N, -1)
+
+        # 计算交集
+        intersection = input_flat * targets_flat
+        N_dice_eff = (2 * intersection.sum(1) + smooth) / (input_flat.sum(1) + targets_flat.sum(1) + smooth)
+        # 计算一个批次中平均每张图的损失
+        loss = 1 - N_dice_eff.sum() / N
+        return loss

Tipsomaly/utils/metrics.py

@@ -0,0 +1,84 @@
+import torch
+import numpy as np
+from skimage import measure
+from torchmetrics import AUROC, AveragePrecision
+from sklearn.metrics import auc, roc_auc_score, average_precision_score, precision_recall_curve
+
+def calc_f1_max(gt, pr):
+    precisions, recalls, _ = precision_recall_curve(gt, pr)
+    f1_scores = (2 * precisions * recalls) / (precisions + recalls)
+    return np.max(f1_scores[np.isfinite(f1_scores)])
+
+
+def cal_pro_score_gpu(device, masks, amaps, max_step=200, expect_fpr=0.3):
+    # GPU implementation using PyTorch
+    if not torch.is_tensor(amaps):
+        amaps = torch.tensor(amaps)
+    amaps = amaps.to(device)
+    masks = masks.to(device)
+    
+    binary_amaps = torch.zeros_like(amaps, dtype=torch.bool, device=device)
+    min_th, max_th = amaps.min().item(), amaps.max().item()
+    delta = (max_th - min_th) / max_step
+    pros, fprs, ths = [], [], []
+    
+    regionprops_list = [measure.regionprops(measure.label(mask.cpu().numpy())) for mask in masks]
+    coords_list = [[(region.coords[:, 0], region.coords[:, 1], len(region.coords)) for region in regionprops] for regionprops in regionprops_list]
+    inverse_masks = 1 - masks
+    tn_pixel = inverse_masks.sum().item() # Pixels that truly has the label of 0
+    for th in torch.arange(min_th, max_th, delta, device=device):
+        binary_amaps[amaps <= th], binary_amaps[amaps > th] = 0, 1
+        pro = []
+        
+        for binary_amap, regions_coords in zip(binary_amaps, coords_list):
+            for coords in regions_coords:
+                tp_pixels = binary_amap[coords[0], coords[1]].sum().item()
+                pro.append(tp_pixels / coords[2])
+        
+        fp_pixels = torch.logical_and(inverse_masks, binary_amaps).sum().item()
+        fpr = fp_pixels / tn_pixel
+        pros.append(np.mean(pro))
+        fprs.append(fpr)
+        ths.append(th.item())
+    
+    pros, fprs, ths = torch.tensor(pros, device=device), torch.tensor(fprs, device=device), torch.tensor(ths, device=device)
+    idxes = fprs < expect_fpr
+    fprs = fprs[idxes]
+    fprs = (fprs - fprs.min()) / (fprs.max() - fprs.min())
+    pro_auc = auc(fprs.cpu().numpy(), pros[idxes].cpu().numpy())
+    return pro_auc
+
+def image_level_metrics(prd, lbl, metric):
+    if len(np.unique(lbl)) < 2:
+        print("only one class present, can not calculate image metrics")
+        return 0
+    
+    if metric == 'auroc':
+        performance = roc_auc_score(lbl, prd)
+    elif metric == 'ap':
+        performance = average_precision_score(lbl, prd)
+    elif metric == 'f1-max':
+        performance = calc_f1_max(lbl, prd)
+    return performance
+
+def pixel_level_metrics(device, prd, lbl, metric):
+    if torch.unique(lbl).numel() < 2:
+        print("only one class present, can not calculate pixel metrics")
+        return 0
+    
+    if metric == 'auroc':
+        performance = AUROC(task="binary")(prd, lbl.to(dtype=torch.long)).item()
+        
+    elif metric == 'aupro':
+        if len(lbl.shape) == 4:
+            lbl = lbl.squeeze(1)
+        if len(prd.shape) == 4:
+            prd = prd.squeeze(1)
+        performance = cal_pro_score_gpu(device, lbl, prd)
+        
+    elif metric == 'ap':     
+        performance = AveragePrecision(task="binary")(prd, lbl.to(dtype=torch.long)).item()
+        
+    elif metric == 'f1-max':
+        performance = calc_f1_max(lbl.cpu().ravel(), prd.cpu().ravel())
+    return performance

Tipsomaly/utils/visualize.py

@@ -0,0 +1,55 @@
+import cv2
+import os
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from PIL import Image
+
+def normalize(pred, max_value=None, min_value=None):
+    if max_value is None or min_value is None:
+        return (pred - pred.min()) / (pred.max() - pred.min())
+    else:
+        return (pred - min_value) / (max_value - min_value)
+
+def apply_ad_scoremap(image, scoremap, alpha=0.5):
+    np_image = np.asarray(image, dtype=float)
+    # Convert scoremap from a PyTorch tensor to a NumPy array
+    if isinstance(scoremap, torch.Tensor):
+        scoremap = scoremap.detach().cpu().numpy()  # Convert tensor to NumPy array
+    scoremap = (scoremap * 255).astype(np.uint8)
+    scoremap = cv2.applyColorMap(scoremap, cv2.COLORMAP_JET)
+    scoremap = cv2.cvtColor(scoremap, cv2.COLOR_BGR2RGB)
+    return (alpha * np_image + (1 - alpha) * scoremap).astype(np.uint8)
+
+def visualizer(pathes, anomaly_map, masks, img_size, cls_name, save_path='./vis_img/', draw_contours=True):
+    for idx, path in enumerate(pathes):
+        cls = path.split('/')[-2]
+        filename = path.split('/')[-1]
+        
+        # Modify filename if contours are enabled
+        if draw_contours:
+            filename_ctr = filename.split('.')[0] + "_cntr." + filename.split('.')[-1]  # Append '_cntr' before file extension
+
+        # Load original image and resize
+        vis = cv2.cvtColor(cv2.resize(cv2.imread(path), (img_size, img_size)), cv2.COLOR_BGR2RGB)
+
+        # Use the provided mask (it's guaranteed to be available)
+        gt_mask = (masks[idx] > 0).astype(np.uint8) * 255  # Convert to binary (0 or 255)
+
+        # Normalize and apply anomaly map
+        mask = normalize(anomaly_map[idx])
+        vis = apply_ad_scoremap(vis, mask)
+
+        # Convert back to BGR for OpenCV
+        vis = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)
+
+        # Save the final visualization
+        save_vis = os.path.join(save_path, str(cls_name), str(cls))
+        os.makedirs(save_vis, exist_ok=True)
+        cv2.imwrite(os.path.join(save_vis, filename), vis)
+        
+        # Find and overlay contours (only if draw_contours is True)
+        if draw_contours:
+            contours, _ = cv2.findContours(gt_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            cv2.drawContours(vis, contours, -1, (120, 251, 120), 2)  # Pale green contours
+            cv2.imwrite(os.path.join(save_vis, filename_ctr), vis)

Tipsomaly/workspaces/trained_on_mvtec_default/vegan-arkansas/checkpoints/args.txt

@@ -0,0 +1,46 @@
+
+Timestamp: 2025-08-24 02:23:07
+image_size: 518
+seed: 111
+metrics: image-pixel-level
+devices: [3]
+model_name: prompt_learning_segonly_518
+sigma: 4
+dataset: ['mvtec']
+dataset_category: 
+class_name: ['all']
+k_shot: 0
+type: train
+log_dir: vegan-arkansas
+image_metrics: ['auroc', 'ap', 'f1-max']
+pixel_metrics: ['auroc', 'aupro', 'f1-max']
+visualize: True
+prompt_learn_method: concat
+cls_seg_los: seg
+data_path: ['/data/alireza/datasets//mvtec/']
+experiment_root: ./workspaces/trained_on_mvtec_prompt_learning_segonly_518/vegan-arkansas
+save_path: ./workspaces/trained_on_mvtec_prompt_learning_segonly_518/vegan-arkansas/checkpoints
+Command arguments: /home/alireza/miniconda3/envs/cuda11_8torch2_5/bin/python /home/alireza/KyotoServer/CLIP-AD/tipsomaly/train.py --model_name prompt_learning_segonly_518 --dataset mvtec --device 3 --cls_seg_los seg
+
+Timestamp: 2025-08-24 02:23:50
+image_size: 518
+seed: 111
+metrics: image-pixel-level
+devices: [3]
+model_name: prompt_learning_segonly_518
+sigma: 4
+dataset: ['mvtec']
+dataset_category: 
+class_name: ['all']
+k_shot: 0
+type: train
+log_dir: vegan-arkansas
+image_metrics: ['auroc', 'ap', 'f1-max']
+pixel_metrics: ['auroc', 'aupro', 'f1-max']
+visualize: True
+prompt_learn_method: concat
+cls_seg_los: seg
+data_path: ['/data/alireza/datasets//mvtec/']
+experiment_root: ./workspaces/trained_on_mvtec_prompt_learning_segonly_518/vegan-arkansas
+save_path: ./workspaces/trained_on_mvtec_prompt_learning_segonly_518/vegan-arkansas/checkpoints
+Command arguments: /home/alireza/miniconda3/envs/cuda11_8torch2_5/bin/python /home/alireza/KyotoServer/CLIP-AD/tipsomaly/train.py --model_name prompt_learning_segonly_518 --dataset mvtec --device 3 --cls_seg_los seg

Tipsomaly/workspaces/trained_on_mvtec_default/vegan-arkansas/log.txt

@@ -0,0 +1,2 @@
+25-08-24 02:27:58.439 - INFO: Epoch 1 Mean Losses: LS CE: 1.6996, LS FC: 0.0422, LS DC P: 0.8976, LS DC N: 0.0594,  Loss: 0.9993
+25-08-24 02:31:43.275 - INFO: Epoch 2 Mean Losses: LS CE: 1.8427, LS FC: 0.0470, LS DC P: 0.7758, LS DC N: 0.0271,  Loss: 0.8499

Tipsomaly/workspaces/trained_on_visa_default/vegan-arkansas/checkpoints/args.txt

@@ -0,0 +1,23 @@
+
+Timestamp: 2025-08-24 02:44:06
+image_size: 518
+seed: 111
+metrics: image-pixel-level
+devices: [3]
+model_name: prompt_learning_segonly_518
+sigma: 4
+dataset: ['visa']
+dataset_category: 
+class_name: ['all']
+k_shot: 0
+type: train
+log_dir: vegan-arkansas
+image_metrics: ['auroc', 'ap', 'f1-max']
+pixel_metrics: ['auroc', 'aupro', 'f1-max']
+visualize: True
+prompt_learn_method: concat
+cls_seg_los: seg
+data_path: ['/data/alireza/datasets//visa/']
+experiment_root: ./workspaces/trained_on_visa_prompt_learning_segonly_518/vegan-arkansas
+save_path: ./workspaces/trained_on_visa_prompt_learning_segonly_518/vegan-arkansas/checkpoints
+Command arguments: /home/alireza/miniconda3/envs/cuda11_8torch2_5/bin/python /home/alireza/KyotoServer/CLIP-AD/tipsomaly/train.py --model_name prompt_learning_segonly_518 --dataset visa --device 3 --cls_seg_los seg

Tipsomaly/workspaces/trained_on_visa_default/vegan-arkansas/log.txt

@@ -0,0 +1,2 @@
+25-08-24 02:49:49.051 - INFO: Epoch 1 Mean Losses: LS CE: 1.6292, LS FC: 0.0135, LS DC P: 0.9782, LS DC N: 0.0234,  Loss: 1.0151
+25-08-24 02:55:10.451 - INFO: Epoch 2 Mean Losses: LS CE: 1.7016, LS FC: 0.0139, LS DC P: 0.9251, LS DC N: 0.0069,  Loss: 0.9459

requirements.txt

@@ -0,0 +1,8 @@
+gradio>=4.44.0
+numpy>=1.26
+pillow>=10.0
+scipy>=1.11
+torch>=2.1
+torchvision>=0.16
+transformers>=4.45
+sentencepiece>=0.1.99