Publish pinned visual components with exact source payloads and native numerical validation

CODE_LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 FoundationVision
+
+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.

LICENSE

@@ -0,0 +1,204 @@
+UniTok model weights are released under Apache-2.0 in the official model card:
+https://huggingface.co/FoundationVision/unitok_tokenizer/tree/7f6ea34f2f669d8d250a4376468c69fe4bcb6c65
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

OPENCLIP_LICENSE

@@ -0,0 +1,24 @@
+MIT License
+
+Copyright (c) 2021 OpenAI
+
+Attribution from the unchanged official vendored open_clip/model.py header.
+Project modifications are also covered by CODE_LICENSE.
+
+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.

README.md

@@ -0,0 +1,40 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: image-to-image
+base_model: FoundationVision/unitok_tokenizer
+tags:
+- unitok
+- image-tokenizer
+- unified-multimodal
+- component-extraction
+---
+# UniTok Large unified visual codec
+
+The complete trained `trainer.unitok` mapping from [the official tokenizer checkpoint](https://huggingface.co/FoundationVision/unitok_tokenizer/tree/7f6ea34f2f669d8d250a4376468c69fe4bcb6c65): ViTamin-L encoder, attention projections, eight normalized codebooks, RGB decoder, semantic image projection and trained CLIP text tower. All 1335 tensors preserve their original names, dtypes and bytes. Discriminator, optimizer and trainer state are excluded.
+
+Install `requirements.txt` before loading. In particular, the native hybrid embedder requires the official `timm==1.0.8`; newer timm constructors use incompatible arguments. The exact text-transformer implementation vendored by the original UniTok project is bundled here; no external OpenCLIP installation is required.
+
+```python
+from huggingface_hub import snapshot_download
+import sys, torch
+folder = snapshot_download('gatilin/unitok_large_visual_codec')
+sys.path.insert(0, folder)
+from component import load_model
+codec = load_model(folder)
+x = torch.zeros(1, 3, 256, 256)  # RGB, NCHW, [-1, 1]
+with torch.inference_mode():
+    ids = codec.img_to_idx(x)  # [B, 8, 256]
+    rgb = codec.idx_to_img(ids)
+    semantic_features = codec.encode_image(x, normalize=True)  # [B, 768]
+```
+
+Each spatial position uses **eight local IDs in [0, 4095]**, one per 4096-by-8 codebook. The checkpoint's `vocab_size=32768` is the sum of the eight codebook sizes, not the upper bound of every local ID. Code lookup normalizes each sub-codebook before concatenating 64-dimensional vectors. The native decoder assumes a square token grid; this release validates its trained 256-by-256 image path.
+
+`preprocessor_config.json` records the deterministic validation transform in the pinned project's `utils/data.py`: resize the shorter edge to 256 with Lanczos, center-crop 256-by-256, convert to RGB, convert to a [0,1] tensor and apply `2*x-1`. Downstream MLLM prompt processors have additional conventions.
+
+Native source is pinned to [FoundationVision/UniTok revision bb8012e](https://github.com/FoundationVision/UniTok/tree/bb8012e8e8d7be66928bedd80506c56a55ea6305), with imports relocated into the bundled `native` package. The text tower directly instantiates the same configured default branch of the original vendored builder, using its unchanged `TextTransformer` and utilities; `native/open_clip/source_model.py` retains the original constructor as a reference. The wrapper disables training checkpointing and strictly restores the complete trained mapping. Text features use `codec.encode_text`; CLIP BPE tokenization should follow the original project. No BPE assets or language generator are included.
+
+`validation.json` records source tensor payload hashes, strict loading, codebook geometry, encode/ID/decode behavior, direct lookup parity, the native quantizer's small floating point difference, and normalized image/text features. These synthetic checks do not reproduce paper quality or downstream accuracy benchmarks. Re-run with `python validate.py`.
+
+Weight terms follow the official model-card Apache-2.0 license (`LICENSE`). Bundled project code follows its MIT license (`CODE_LICENSE`); `OPENCLIP_LICENSE` and the unchanged reference source retain the original OpenAI attribution. Original third-party notices in individual files are retained. Source container integrity, mapping, per-tensor hashes and code provenance are included.

code_provenance.json

@@ -0,0 +1,66 @@
+{
+  "native/models/unitok.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/models/unitok.py",
+    "original_sha256": "27872719d35827715e478556e857efa1236b6575219db2aac8206ae354e04972",
+    "modifications": [
+      "Package imports relocated into native namespace.",
+      "Instantiate the exact configured non-HF default branch of the official vendored _build_text_tower directly with native.open_clip.transformer.TextTransformer; its original constructor remains in native/open_clip/source_model.py."
+    ]
+  },
+  "native/models/quant.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/models/quant.py",
+    "original_sha256": "9d121a0d8d3d52127afd118df3c119a30a4dea784433ea6403a3be299945375a",
+    "modifications": [
+      "Package imports relocated into native namespace."
+    ]
+  },
+  "native/models/vitamin.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/models/vitamin.py",
+    "original_sha256": "9c2df37cbad670786a4ac3a5eeece4fe3ab3de57c9c5848b4f6bdd44f76d4600",
+    "modifications": [
+      "Package imports relocated into native namespace."
+    ]
+  },
+  "native/models/vqvae.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/models/vqvae.py",
+    "original_sha256": "9ddcb17ea31936e80bdf7cd749082314aa7a6686c0c60e79e64edc467f75323d",
+    "modifications": [
+      "Package imports relocated into native namespace."
+    ]
+  },
+  "native/utils/config.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/utils/config.py",
+    "original_sha256": "94b0c9637ef1a8dfd710c42b97ad4e8175f56573a3705b912b217b6b4b527037",
+    "modifications": [
+      "Package imports relocated into native namespace."
+    ]
+  },
+  "native/utils/dist.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/utils/dist.py",
+    "original_sha256": "4ad98f72e600a666795eb074329f4e6cfd8c182776fb7f9e4d9d39ae2dd83524",
+    "modifications": []
+  },
+  "CODE_LICENSE": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/LICENSE",
+    "original_sha256": "a318f4da9d043e05f535b65774febde7572908c97562cc14e49a2f3687f5a453",
+    "modifications": []
+  },
+  "native/open_clip/transformer.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/open_clip/transformer.py",
+    "original_sha256": "dfaef403435d5766aeea8286fca0f9fc46285ee25fd28e87ba197bf51f1f676e",
+    "modifications": [],
+    "purpose": "Exact official vendored text computation; no external OpenCLIP dependency."
+  },
+  "native/open_clip/utils.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/open_clip/utils.py",
+    "original_sha256": "f0267bb9d5402b1b5b38a90f96f6cd9002ac21d2db0f3f3fbee8ed48b2f46047",
+    "modifications": [],
+    "purpose": "Exact official vendored text computation; no external OpenCLIP dependency."
+  },
+  "native/open_clip/source_model.py": {
+    "url": "https://raw.githubusercontent.com/FoundationVision/UniTok/bb8012e8e8d7be66928bedd80506c56a55ea6305/open_clip/model.py",
+    "original_sha256": "2a7bc461f7cfe70497e7ce8541f75a8b67dcbd3d93685cfd8671ebf98d337150",
+    "modifications": [],
+    "purpose": "Original constructor reference only."
+  }
+}

component.py

@@ -0,0 +1,17 @@
+"""Strict loader preserving the complete trained UniTok state, including semantic heads."""
+import json
+from pathlib import Path
+from types import SimpleNamespace
+import torch
+from safetensors.torch import load_file
+from native.models.unitok import UniTok
+
+
+def load_model(path=None, dtype=torch.float32, device='cpu'):
+    root = Path(path or Path(__file__).parent)
+    cfg = json.loads((root / 'config.json').read_text())
+    cfg.update(device='cpu', grad_ckpt=False)
+    # The native hybrid embedder runs a CPU shape probe in its constructor.
+    model = UniTok(SimpleNamespace(**cfg))
+    model.load_state_dict(load_file(root / 'model.safetensors'), strict=True, assign=True)
+    return model.to(device=device, dtype=dtype).eval()

config.json

@@ -0,0 +1,116 @@
+{
+  "model": "vitamin_large",
+  "exp_name": "unitok_large_causal",
+  "output_dir": "local_output",
+  "resume_from": "",
+  "lpips_path": "",
+  "dino_path": "",
+  "fid_eval_src": "",
+  "fid_eval_dst": "",
+  "vis_img_dir": "asset/vis_imgs/",
+  "fid_feature_extractor": "",
+  "clip_pretrain_path": "",
+  "fp16": false,
+  "bf16": true,
+  "tf32": true,
+  "compile_model": false,
+  "ddp_static": false,
+  "grad_ckpt": true,
+  "grad_accu": 1,
+  "dtype": "torch.bfloat16",
+  "train_data": "",
+  "val_data": null,
+  "dataset_type": "webdataset",
+  "imagenet_val": "",
+  "imagenet_v2": null,
+  "subset_ratio": 1.0,
+  "img_size": 256,
+  "resize_ratio": 1.125,
+  "hflip": false,
+  "workers": 16,
+  "train_num_samples": 1280000000,
+  "train_data_upsampling_factors": null,
+  "dataset_resampled": false,
+  "use_aug": false,
+  "vocab_size": 32768,
+  "vocab_width": 64,
+  "vocab_norm": true,
+  "vq_beta": 0.25,
+  "num_codebooks": 8,
+  "quant_proj": "attn",
+  "embed_dim": 768,
+  "num_query": 0,
+  "use_clip_pretrain": false,
+  "patch_size": 16,
+  "drop_path": 0.1,
+  "text_width": 768,
+  "text_heads": 12,
+  "text_layers": 12,
+  "text_vocab_size": 49408,
+  "text_context_length": 77,
+  "local_loss": true,
+  "gather_with_grad": true,
+  "pretrained_clip": null,
+  "pretrained_clip_text": null,
+  "lock_text": false,
+  "lock_text_unlocked_layers": 0,
+  "lock_text_freeze_layer_norm": false,
+  "force_custom_text": false,
+  "force_custom_vision": false,
+  "zeroshot_eval_freq": 1,
+  "dino_depth": 12,
+  "dino_kernel_size": 9,
+  "disc_norm": "gn",
+  "disc_aug_prob": 1.0,
+  "disc_specnorm": false,
+  "step_disc_every": 1,
+  "vae_init": -0.5,
+  "vocab_init": -1,
+  "disc_init": -0.5,
+  "epoch": 1,
+  "local_bs": 56,
+  "vae_local_bs": 56,
+  "global_bs": 16384,
+  "lr": 0.0005,
+  "wd": 0.02,
+  "disc_lr": 2e-05,
+  "disc_wd": 0.2,
+  "grad_clip": 10,
+  "ema": 0.9999,
+  "warmup_iter": null,
+  "warmup_ep": 0.01,
+  "disc_start_ep": 0.375,
+  "disc_warmup_ep": 0.03,
+  "schedule": "cos",
+  "lr_start_ratio": 0.0,
+  "lr_end_ratio": 0.1,
+  "disc_lr_end_ratio": 0.1,
+  "custom_lr_multiplier": null,
+  "optimizer": "adamw",
+  "optim_eps": 1e-06,
+  "fuse_opt": false,
+  "optim_beta": "0.9_0.95",
+  "disc_optim_beta": "0.5_0.9",
+  "l1": 0.2,
+  "l2": 1.0,
+  "lp": 1.0,
+  "lpr": 48,
+  "ld": 0.4,
+  "le": 0.0,
+  "lq": 1.0,
+  "lc": 1.0,
+  "e_temp": 0.01,
+  "gada": 1,
+  "bcr": 4.0,
+  "bcr_cut": 0.2,
+  "dcrit": "hg",
+  "report_wandb": true,
+  "wandb_notes": null,
+  "run_id": null,
+  "eval_per_epoch": 8,
+  "dbg_unused_param": false,
+  "dbg_nan": false,
+  "seed": null,
+  "deterministic": false,
+  "same_seed_for_all_ranks": 0
+}

extraction.json

@@ -0,0 +1,15 @@
+{
+  "source_repo": "FoundationVision/unitok_tokenizer",
+  "source_revision": "7f6ea34f2f669d8d250a4376468c69fe4bcb6c65",
+  "source_file": "unitok_tokenizer.pth",
+  "prefixes": [],
+  "source_tensor_bytes_equal": true,
+  "sha256": "364bbd3e3f453f19dba447805f347aabdd7abdc7da20186364f5b2961c541b5c",
+  "bytes": 3281477448,
+  "verification": "Every HTTP range received at its requested exact offsets and persisted-byte SHA-256 checked; whole original file also matches official LFS SHA-256.",
+  "source_file_lfs_sha256": "b1e39e3df707c909c81de0affcbf45075d40d9f156ef0e200b037bcc9b6db81a",
+  "source_field": "trainer.unitok",
+  "tensor_count": 1335,
+  "source_container_sha256": "b1e39e3df707c909c81de0affcbf45075d40d9f156ef0e200b037bcc9b6db81a",
+  "transform": "Extract trainer.unitok only; retain trained encoder, quantizer, decoder, semantic projection and text tower. Discriminator/training state excluded; original names, dtypes and bytes preserved."
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:364bbd3e3f453f19dba447805f347aabdd7abdc7da20186364f5b2961c541b5c
+size 3281477448

native/models/quant.py

@@ -0,0 +1,185 @@
+import torch
+from typing import List, Tuple
+from torch.nn import functional as F
+from torch import distributed as tdist, nn as nn
+
+from native.utils import dist
+
+
+def get_entropy_loss(latent_embed, codebook_embed, inv_entropy_tau):
+    E_dist = latent_embed.square().sum(dim=1, keepdim=True) + codebook_embed.square().sum(dim=1, keepdim=False)
+    E_dist.addmm_(latent_embed, codebook_embed.T, alpha=-2, beta=1)  # E_dist: (N, vocab_size)
+    logits = -E_dist.float().mul_(inv_entropy_tau)
+    # calc per_sample_entropy
+    prob, log_prob = logits.softmax(dim=-1), logits.log_softmax(dim=-1)  # both are (N, vocab_size)
+    per_sample_entropy = torch.mean((-prob * log_prob).sum(dim=-1))
+    # calc codebook_entropy
+    avg_prob = prob.mean(dim=0)  # (vocab_size,)
+    log_avg_prob = torch.log(avg_prob + 1e-7)
+    codebook_entropy = (-avg_prob * log_avg_prob).sum()
+    # calc entropy_loss
+    entropy_loss = per_sample_entropy - codebook_entropy
+    return entropy_loss
+
+
+class NormalizedEmbedding(nn.Embedding):
+    def __init__(self, num_embeddings: int, embedding_dim: int):
+        super().__init__(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
+        # self.norm_scale = nn.Parameter(torch.tensor(0.0, dtype=torch.float32))
+
+    def forward(self, idx):
+        return F.embedding(
+            idx, F.normalize(self.weight, dim=1), self.padding_idx, self.max_norm,
+            self.norm_type, self.scale_grad_by_freq, self.sparse
+        )
+
+    def get_norm_weight(self):
+        return F.normalize(self.weight, dim=1)
+
+
+class ResConv(nn.Conv2d):
+    def __init__(self, embed_dim, quant_resi):
+        ks = 3 if quant_resi < 0 else 1
+        super().__init__(in_channels=embed_dim, out_channels=embed_dim, kernel_size=ks, stride=1, padding=ks // 2)
+        self.resi_ratio = abs(quant_resi)
+
+    def forward(self, h_BChw):
+        return h_BChw.mul(1 - self.resi_ratio) + super().forward(h_BChw).mul_(self.resi_ratio)
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        vocab_width: int,
+        beta: float = 0.25,
+        use_entropy_loss=False,
+        entropy_temp=0.01,
+    ):
+        super().__init__()
+        self.beta = beta
+        self.vocab_size = vocab_size
+        self.vocab_width = vocab_width
+        self.vocab_usage_record_times: int = 0
+        self.register_buffer('vocab_usage', torch.zeros(self.vocab_size))
+        self.codebook = NormalizedEmbedding(self.vocab_size, self.vocab_width)
+
+        self.use_entropy_loss = use_entropy_loss
+        self.inv_entropy_tau = 1 / entropy_temp
+
+    def init_vocab(self, eini: float):
+        if eini > 0:
+            nn.init.trunc_normal_(self.codebook.weight.data, std=eini)
+        elif eini < 0:
+            base = self.vocab_width ** -0.5
+            base /= 36
+            self.codebook.weight.data.uniform_(-abs(eini) * base, abs(eini) * base)
+
+    def extra_repr(self) -> str:
+        return f'beta={self.beta:g}'
+
+    def forward(self, features):
+        B, L, C = features.shape
+        features = features.reshape(-1, C)
+        features = F.normalize(features, dim=-1).float()
+        codebook_embed = self.codebook.get_norm_weight()
+        indices = torch.argmax(features.detach() @ codebook_embed.T, dim=1)
+        entropy_loss = get_entropy_loss(features, codebook_embed, self.inv_entropy_tau) if self.use_entropy_loss else 0
+        features_hat = self.codebook(indices)
+
+        # calc loss
+        vq_loss = F.mse_loss(features_hat.detach(), features).mul_(self.beta) + F.mse_loss(features_hat,
+                                                                                           features.detach())
+        features_hat = (features_hat.detach() - features.detach()).add_(features)
+
+        # update vocab_usage
+        prob_per_class_is_chosen = indices.bincount(minlength=self.vocab_size).float()
+        handler = tdist.all_reduce(prob_per_class_is_chosen, async_op=True) if (
+                self.training and dist.initialized()) else None
+        if handler is not None:
+            handler.wait()
+        prob_per_class_is_chosen /= prob_per_class_is_chosen.sum()
+        vocab_usage = (prob_per_class_is_chosen > 0.01 / self.vocab_size).float().mean().mul_(100)
+        if self.vocab_usage_record_times == 0:
+            self.vocab_usage.copy_(prob_per_class_is_chosen)
+        elif self.vocab_usage_record_times < 100:
+            self.vocab_usage.mul_(0.9).add_(prob_per_class_is_chosen, alpha=0.1)
+        else:
+            self.vocab_usage.mul_(0.99).add_(prob_per_class_is_chosen, alpha=0.01)
+        self.vocab_usage_record_times += 1
+
+        return features_hat.view(B, L, C), vq_loss, entropy_loss, vocab_usage
+
+    def f_to_idx(self, features):
+        B, L, C = features.shape
+        features = features.reshape(-1, C)
+        features = F.normalize(features, dim=-1).float()
+        codebook_embed = self.codebook.get_norm_weight().float()
+        indices = torch.argmax(features.detach() @ codebook_embed.T, dim=1)
+        return indices.view(B, L)
+
+
+class VectorQuantizerM(nn.Module):
+    def __init__(
+        self,
+        vocab_size,
+        vocab_width,
+        beta=0.25,
+        use_entropy_loss=False,
+        entropy_temp=0.01,
+        num_codebooks=16
+    ):
+        super().__init__()
+        self.num_codebooks = num_codebooks
+        self.codebooks = nn.ModuleList()
+        for _ in range(num_codebooks):
+            codebook = VectorQuantizer(
+                vocab_size=vocab_size // num_codebooks,
+                vocab_width=vocab_width // num_codebooks,
+                beta=beta,
+                use_entropy_loss=use_entropy_loss,
+                entropy_temp=entropy_temp,
+            )
+            self.codebooks.append(codebook)
+
+    def init_vocab(self, eini: float):
+        for codebook in self.codebooks:
+            codebook.init_vocab(eini)
+
+    def f_to_idx(self, features):
+        indices = []
+        chunk_size = features.shape[-1] // self.num_codebooks
+        splited_features = features.split(chunk_size, dim=-1)
+        for i, codebook in enumerate(self.codebooks):
+            indices.append(codebook.f_to_idx(splited_features[i]))
+        indices = torch.stack(indices, dim=1)
+        return indices
+
+    def idx_to_f(self, indices):
+        assert indices.shape[1] == self.num_codebooks
+        latent_features = []
+        for i, codebook in enumerate(self.codebooks):
+            sub_indices = indices[:, i].flatten(start_dim=1)
+            latent_feature = codebook.codebook(sub_indices)
+            latent_features.append(latent_feature)
+        latent_features = torch.cat(latent_features, dim=-1)
+        return latent_features
+
+    def forward(self, features):
+        latent_features = []
+        global_vq_loss = 0.
+        global_entropy_loss = 0.
+        global_vocab_usage = 0.
+        chunk_size = features.shape[-1] // self.num_codebooks
+        splited_features = features.split(chunk_size, dim=-1)
+        for i, codebook in enumerate(self.codebooks):
+            latent_feature, vq_loss, entropy_loss, vocab_usage = codebook(splited_features[i])
+            latent_features.append(latent_feature)
+            global_vq_loss += vq_loss
+            global_entropy_loss += entropy_loss
+            global_vocab_usage += vocab_usage
+        latent_features = torch.cat(latent_features, dim=-1)
+        global_entropy_loss /= self.num_codebooks
+        global_vq_loss /= self.num_codebooks
+        global_vocab_usage /= self.num_codebooks
+        return latent_features, global_vq_loss, global_entropy_loss, global_vocab_usage

native/models/unitok.py

@@ -0,0 +1,185 @@
+import timm
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from contextlib import nullcontext
+
+from native.models.vitamin import GeGluMlp, ViTaminDecoder
+from native.models.quant import VectorQuantizerM
+from native.models.vqvae import AttnProjection
+
+
+class UniTok(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        self.num_query = args.num_query
+
+        self.encoder = timm.create_model(
+            args.model,
+            patch_size=1,
+            fc_norm=False,
+            drop_rate=0.0,
+            num_classes=0,
+            global_pool='',
+            pos_embed='none',
+            class_token=False,
+            mlp_layer=GeGluMlp,
+            reg_tokens=args.num_query,
+            img_size=args.img_size,
+            drop_path_rate=args.drop_path,
+        )
+        self.encoder.pos_embed = nn.Parameter(torch.zeros(1, 1, self.encoder.embed_dim), requires_grad=False)
+
+        if args.quant_proj == 'linear':
+            self.quant_proj = nn.Linear(self.encoder.embed_dim, args.vocab_width)
+        elif args.quant_proj == 'attn':
+            self.quant_proj = AttnProjection(self.encoder.embed_dim, args.vocab_width, self.encoder.embed_dim // args.vocab_width)
+        else:
+            raise NotImplementedError
+
+        self.quantizer = VectorQuantizerM(
+            vocab_size=args.vocab_size,
+            vocab_width=args.vocab_width,
+            beta=args.vq_beta,
+            use_entropy_loss=args.le > 0,
+            entropy_temp=args.e_temp,
+            num_codebooks=args.num_codebooks,
+        )
+
+        if args.quant_proj == 'linear':
+            self.post_quant_proj = nn.Linear(args.vocab_width, self.encoder.embed_dim)
+        elif args.quant_proj == 'attn':
+            self.post_quant_proj = AttnProjection(args.vocab_width, self.encoder.embed_dim, self.encoder.embed_dim // args.vocab_width)
+        else:
+            raise NotImplementedError
+
+        self.decoder = ViTaminDecoder(
+            args.model,
+            num_query=args.num_query,
+            img_size=args.img_size,
+            drop_path=args.drop_path,
+            grad_ckpt=args.grad_ckpt,
+        )
+
+        text_cfg = {
+            "width": args.text_width,
+            "heads": args.text_heads,
+            "layers": args.text_layers,
+            "vocab_size": args.text_vocab_size,
+            "context_length": args.text_context_length,
+        }
+        # Exact non-HF default branch of the pinned project's vendored builder.
+        from native.open_clip.transformer import TextTransformer
+        self.text_encoder = TextTransformer(output_dim=args.embed_dim, **text_cfg)
+
+        self.fc_norm = nn.LayerNorm(self.encoder.embed_dim, eps=1e-6)
+        self.projection = nn.Linear(self.encoder.embed_dim, args.embed_dim)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.context_length = self.text_encoder.context_length
+        self.vocab_size = self.text_encoder.vocab_size
+        self.maybe_record_function = nullcontext
+
+        self.text_no_grad = False
+        self.encoder.set_grad_checkpointing(args.grad_ckpt)
+        self.text_encoder.set_grad_checkpointing(args.grad_ckpt)
+
+    def forward(self, img, vae_bs, text=None, ret_usages=False):
+        img_tokens = self.encoder(img).float()
+        with torch.cuda.amp.autocast(enabled=False):
+            img_tokens = torch.utils.checkpoint.checkpoint(self.quant_proj, img_tokens, use_reentrant=False)
+            img_tokens, vq_loss, entropy_loss, usages = self.quantizer(img_tokens)
+            img_tokens = torch.utils.checkpoint.checkpoint(self.post_quant_proj, img_tokens, use_reentrant=False)
+        img_rec = self.decoder(img_tokens[:vae_bs]).float()
+
+        clip_visual = img_tokens.mean(dim=1)
+        clip_visual = self.projection(self.fc_norm(clip_visual))
+        clip_visual = F.normalize(clip_visual, dim=-1)
+        if text is not None:
+            clip_text = self.text_encoder(text)
+            clip_text = F.normalize(clip_text, dim=-1)
+        else:
+            clip_text = None
+
+        output_dict = {
+            "img_rec": img_rec,
+            "vq_loss": vq_loss,
+            "entropy_loss": entropy_loss,
+            "codebook_usages": usages,
+            "clip_image_features": clip_visual,
+            "clip_text_features": clip_text,
+            "logit_scale": self.logit_scale.exp()
+        }
+        return output_dict
+
+    def encode_image(self, image, normalize: bool = False):
+        img_tokens = self.encoder(image)
+        img_tokens = self.quant_proj(img_tokens)
+        img_indices = self.quantizer.f_to_idx(img_tokens)
+        img_tokens = self.quantizer.idx_to_f(img_indices)
+        img_tokens = self.post_quant_proj(img_tokens)
+        features = img_tokens.mean(dim=1)
+        features = self.projection(self.fc_norm(features))
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text_encoder(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def img_to_idx(self, img):
+        features = self.encoder(img).float()
+        features = self.quant_proj(features)
+        return self.quantizer.f_to_idx(features)
+
+    def idx_to_img(self, indices):
+        features = self.quantizer.idx_to_f(indices)
+        features = self.post_quant_proj(features)
+        img = self.decoder(features).clamp_(-1, 1)
+        return img
+
+    def img_to_reconstructed_img(self, image) -> torch.Tensor:
+        img_tokens = self.encoder(image)
+        img_tokens = self.quant_proj(img_tokens)
+        img_tokens, _, _, _ = self.quantizer(img_tokens)
+        img_tokens = self.post_quant_proj(img_tokens)
+        img_rec = self.decoder(img_tokens).clamp_(-1, 1)
+        return img_rec
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True, unlock_text_proj=False):
+        self.text.lock(unlocked_layers, freeze_layer_norm, unlock_text_proj)
+        self.text_no_grad = True
+
+
+if __name__ == '__main__':
+    model = timm.create_model(
+        'vitamin_base',
+        patch_size=1,
+        fc_norm=True,
+        drop_rate=0.0,
+        num_classes=0,
+        global_pool='',
+        pos_embed='none',
+        class_token=False,
+        mlp_layer=GeGluMlp,
+        reg_tokens=0,
+        img_size=256,
+        drop_path_rate=0.1,
+    )
+    model.pos_embed = nn.Parameter(torch.zeros(1, 1, model.embed_dim), requires_grad=False)
+
+    model_dict = model.state_dict()
+    ckpt_dict = torch.load('ViTamin-B/pytorch_model.bin')
+    visual_dict = dict()
+    for k, v in ckpt_dict.items():
+        if k.startswith('visual.'):
+            if 'head' in k or 'pos_embed' in k:
+                continue
+            new_k = k.replace('visual.trunk.', '')
+            visual_dict[new_k] = v
+
+    model.load_state_dict(visual_dict, strict=False)
+    print(set(model_dict.keys()) - set(visual_dict.keys()))
+    print(set(visual_dict.keys() - set(model_dict.keys())))
+

native/models/vitamin.py

@@ -0,0 +1,793 @@
+"""
+TODO: FIXME:
+/usr/local/lib/python3.9/dist-packages/torch/autograd/__init__.py:251: UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance.
+grad.sizes() = [256, 1024, 1, 1], strides() = [1024, 1, 1024, 1024]
+bucket_view.sizes() = [256, 1024, 1, 1], strides() = [1024, 1, 1, 1] (Triggered internally at ../torch/csrc/distributed/c10d/reducer.cpp:334.)  Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
+/usr/local/lib/python3.9/dist-packages/torch/autograd/__init__.py:251: UserWarning: Grad strides do not match bucket view strides. This may indicate grad was not created according to the gradient layout contract, or that the param's strides changed since DDP was constructed.  This is not an error, but may impair performance.
+grad.sizes() = [256, 1024, 1, 1], strides() = [1024, 1, 1024, 1024]
+
+"""
+
+""" ViTamin
+
+Paper: Designing Scalable Vison Models in the Vision-Language Era
+
+@misc{chen2023designing,
+      title={Designing Scalable Vison Models in the Vision-Language Era},
+      author={Jieneng Chen and Qihang Yu and Xiaohui Shen and Alan Yuille and Liang-Cheih Chen},
+      year={2023},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+
+Based on Apache 2.0 licensed code at https://github.com/ViTamin/ViTamin
+
+Modifications and timm support by Jieneng Chen 2023
+
+Reference:
+https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py
+https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer_hybrid.py
+"""
+
+import math
+from dataclasses import dataclass
+from functools import partial
+import torch.nn.functional as F
+from typing import Optional, Tuple, Union
+
+import timm
+import torch
+import torch.nn as nn
+from timm.layers import to_2tuple
+from timm.layers.norm_act import _create_act
+from timm.models._builder import build_model_with_cfg
+from timm.models._manipulate import checkpoint_seq, named_apply
+from timm.models._registry import register_model
+from timm.models.layers import DropPath
+from timm.models.layers import create_conv2d, get_norm_act_layer, get_norm_layer, make_divisible
+from timm.models.vision_transformer import VisionTransformer, checkpoint_filter_fn
+from timm.models.vision_transformer_hybrid import HybridEmbed
+from torch.utils.checkpoint import checkpoint
+
+from native.utils.config import Args
+DropPath.__repr__ = lambda self: f'{type(self).__name__}(...)'
+
+
+@dataclass
+class VitConvCfg:
+    expand_ratio: float = 4.0
+    expand_output: bool = True  # calculate expansion channels from output (vs input chs)
+    kernel_size: int = 3
+    group_size: int = 1  # 1 == depthwise
+    pre_norm_act: bool = False  # activation after pre-norm
+    stride_mode: str = 'dw'  # stride done via one of 'pool', '1x1', 'dw'
+    pool_type: str = 'avg2'
+    downsample_pool_type: str = 'avg2'
+    act_layer: str = 'gelu' # stem & stage 1234
+    act_layer1: str = 'gelu' # stage 1234
+    act_layer2: str = 'gelu' # stage 1234
+    norm_layer: str = ''
+    norm_layer_cl: str = ''
+    norm_eps: Optional[float] = None
+    down_shortcut: Optional[bool] = True
+    mlp: str = 'mlp'
+
+    def __post_init__(self):
+        # mbconv vs convnext blocks have different defaults, set in post_init to avoid explicit config args
+        use_mbconv = True
+        if not self.norm_layer:
+            self.norm_layer = 'batchnorm2d' if use_mbconv else 'layernorm2d'
+        if not self.norm_layer_cl and not use_mbconv:
+            self.norm_layer_cl = 'layernorm'
+        if self.norm_eps is None:
+            self.norm_eps = 1e-5 if use_mbconv else 1e-6
+        self.downsample_pool_type = self.downsample_pool_type or self.pool_type
+
+
+@dataclass
+class VitCfg:
+    # embed_dim: Tuple[int, ...] = (96, 192, 384, 768)
+    embed_dim: Tuple[Union[int, Tuple[int, ...]], ...] = (96, 192, 384, 768)
+    depths: Tuple[Union[int, Tuple[int, ...]], ...] = (2, 3, 5, 2)
+    stem_width: int = 64
+    conv_cfg: VitConvCfg = None
+    weight_init: str = 'vit_eff'
+    head_type: str = ""
+    stem_type: str = "stem"
+    ln2d_permute: bool = True
+    # memory_format: str=""
+
+
+def _init_conv(module, name, scheme=''):
+    if isinstance(module, nn.Conv2d):
+        fan_out = module.kernel_size[0] * module.kernel_size[1] * module.out_channels
+        fan_out //= module.groups
+        nn.init.normal_(module.weight, 0, math.sqrt(2.0 / fan_out))
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+class Stem(nn.Module):
+    def __init__(
+            self,
+            in_chs: int,
+            out_chs: int,
+            act_layer: str = 'gelu',
+            norm_layer: str = 'layernorm2d',
+            norm_eps: float = 1e-6,
+            bias: bool = True,
+    ):
+        super().__init__()
+        self.grad_checkpointing=False
+        norm_act_layer = partial(get_norm_act_layer(norm_layer, act_layer), eps=norm_eps)
+        self.out_chs = out_chs
+        self.conv1 = create_conv2d(in_chs, out_chs, 3, stride=2, bias=bias)
+        self.norm1 = norm_act_layer(out_chs)
+        self.conv2 = create_conv2d(out_chs, out_chs, 3, stride=1, bias=bias)
+        named_apply(_init_conv, self)
+
+    def forward(self, x):
+        if self.grad_checkpointing:
+            x = checkpoint(self.conv1, x)
+            x = self.norm1(x)
+            x = checkpoint(self.conv2, x)
+        else:
+            x = self.conv1(x)
+            x = self.norm1(x)
+            x = self.conv2(x)
+
+        return x
+
+
+class Downsample2d(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            dim_out: int,
+            pool_type: str = 'avg2',
+            bias: bool = True,
+    ):
+        super().__init__()
+        self.pool = nn.AvgPool2d(kernel_size=3, stride=2, padding=1, count_include_pad=False)
+
+        if dim != dim_out:
+            self.expand = nn.Conv2d(dim, dim_out, 1, bias=bias) # 1x1 conv
+        else:
+            self.expand = nn.Identity()
+
+    def forward(self, x):
+        x = self.pool(x)  # spatial downsample
+        x = self.expand(x)  # expand chs
+        return x
+
+
+class StridedConv(nn.Module):
+    """ downsample 2d as well
+    """
+    def __init__(
+            self,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            in_chans=3,
+            embed_dim=768,
+            ln2d_permute=True
+    ):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)
+        self.permute = ln2d_permute # TODO: disable
+        norm_layer = partial(get_norm_layer('layernorm2d'), eps=1e-6)
+        self.norm = norm_layer(in_chans) # affine over C
+
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.proj(x)
+        return x
+
+
+class MbConvLNBlock(nn.Module):
+    """ Pre-Norm Conv Block - 1x1 - kxk - 1x1, w/ inverted bottleneck (expand)
+    """
+    def __init__(
+        self,
+        in_chs: int,
+        out_chs: int,
+        stride: int = 1,
+        drop_path: float = 0.,
+        kernel_size: int = 3,
+        norm_layer: str = 'layernorm2d',
+        norm_eps: float = 1e-6,
+        act_layer: str = 'gelu',
+        expand_ratio: float = 4.0,
+    ):
+        super(MbConvLNBlock, self).__init__()
+        self.stride, self.in_chs, self.out_chs = stride, in_chs, out_chs
+        mid_chs = make_divisible(out_chs * expand_ratio)
+        prenorm_act_layer = partial(get_norm_act_layer(norm_layer, act_layer), eps=norm_eps)
+
+        if stride == 2:
+            self.shortcut = Downsample2d(in_chs, out_chs, pool_type='avg', bias=True)
+        elif in_chs != out_chs:
+            self.shortcut = nn.Conv2d(in_chs, out_chs, 1, bias=True)
+        else:
+            self.shortcut = nn.Identity()
+
+        self.pre_norm = prenorm_act_layer(in_chs, apply_act=False)
+        self.down = nn.Identity()
+        self.conv1_1x1 = create_conv2d(in_chs, mid_chs, 1, stride=1, bias=True)
+        self.act1 = _create_act(act_layer, inplace=True)
+        self.act2 = _create_act(act_layer, inplace=True)
+
+        self.conv2_kxk = create_conv2d(mid_chs, mid_chs, kernel_size, stride=stride, dilation=1, groups=mid_chs, bias=True)
+        self.conv3_1x1 = create_conv2d(mid_chs, out_chs, 1, bias=True)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def init_weights(self, scheme=''):
+        named_apply(partial(_init_conv, scheme=scheme), self)
+
+    def forward(self, x):
+        shortcut = self.shortcut(x)
+
+        x = self.pre_norm(x)
+        x = self.down(x) # nn.Identity()
+
+        # 1x1 expansion conv & act
+        x = self.conv1_1x1(x)
+        x = self.act1(x)
+
+        # (strided) depthwise 3x3 conv & act
+        x = self.conv2_kxk(x)
+        x = self.act2(x)
+
+        # 1x1 linear projection to output width
+        x = self.conv3_1x1(x)
+        x = self.drop_path(x) + shortcut
+
+        return x
+
+
+class MbConvStages(nn.Module):
+    """ MobileConv for stage 1 and stage 2 of ViTamin
+    """
+    def __init__(
+            self,
+        cfg: VitCfg,
+        img_size: Union[int, Tuple[int, int]] = 224, # place holder
+        in_chans: int = 3,
+    ):
+        super().__init__()
+        self.grad_checkpointing = False
+        self.stem = Stem(
+            in_chs=in_chans,
+            out_chs=cfg.stem_width,
+        )
+        stages = []
+        self.num_stages = len(cfg.embed_dim)
+        for s, dim in enumerate(cfg.embed_dim[:2]): # stage
+            blocks = []
+            stage_in_chs = cfg.embed_dim[s-1] if s>0 else cfg.stem_width
+            for d in range(cfg.depths[s]):
+                blocks += [MbConvLNBlock(
+                        in_chs = stage_in_chs if d==0 else dim,
+                        out_chs = dim,
+                        stride = 2 if d == 0 else 1,
+                        # cfg = cfg.conv_cfg,
+                    )]
+            blocks = nn.Sequential(*blocks)
+            stages += [blocks]
+
+        self.stages = nn.ModuleList(stages)
+        self.pool = StridedConv(
+                        stride=2,
+                        in_chans=cfg.embed_dim[1],
+                        embed_dim=cfg.embed_dim[2]
+                    )
+
+    def forward(self, x):
+        x = self.stem(x)
+        if self.grad_checkpointing and not torch.jit.is_scripting():
+            for stage in self.stages:
+                x = checkpoint_seq(stage, x)
+            x = checkpoint(self.pool, x)
+        else:
+            for stage in self.stages:
+                x = stage(x)
+            x = self.pool(x)
+
+        return x
+
+
+class GeGluMlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features,
+        act_layer = None,
+        drop = 0.0,
+    ):
+        super().__init__()
+        norm_layer = partial(get_norm_layer('layernorm'), eps=1e-6)
+        self.norm = norm_layer(in_features)
+        self.act = nn.GELU(approximate='tanh')
+        self.w0 = nn.Linear(in_features, hidden_features)
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(hidden_features, in_features)
+
+    def forward(self, x):
+        x = self.norm(x)
+        x = self.act(self.w0(x)) * self.w1(x)
+        x = self.w2(x)
+        return x
+
+
+class HybridEmbed(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+    def __init__(
+        self,
+        backbone,
+        img_size=256,
+        patch_size=1,
+        feature_size=None,
+        in_chans=3,
+        embed_dim=1024,
+        bias=True,
+        dynamic_img_pad=False,
+    ):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.backbone = backbone
+        with torch.no_grad():
+            training = backbone.training
+            if training:
+                backbone.eval()
+            o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))
+            if isinstance(o, (list, tuple)):
+                o = o[-1]  # last feature if backbone outputs list/tuple of features
+            feature_size = o.shape[-2:]
+            feature_dim = o.shape[1]
+            backbone.train(training)
+
+        assert feature_size[0] % patch_size[0] == 0 and feature_size[1] % patch_size[1] == 0
+        self.grid_size = (feature_size[0] // patch_size[0], feature_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.proj = nn.Identity()
+
+    def forward(self, x):
+        x = self.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)
+        return x
+
+
+class Upsample2d(nn.Module):
+    def __init__(self, dim, dim_out):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(dim, dim_out, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv(x)
+        return x
+
+
+class InvMbConvLNBlock(nn.Module):
+    """ Pre-Norm Conv Block - 1x1 - kxk - 1x1, w/ inverted bottleneck (expand)
+    """
+    def __init__(
+        self,
+        in_chs: int,
+        out_chs: int,
+        stride: int = 1,
+        drop_path: float = 0.,
+        kernel_size: int = 3,
+        norm_layer: str = 'layernorm2d',
+        norm_eps: float = 1e-6,
+        act_layer: str = 'gelu',
+        expand_ratio: float = 4.0,
+    ):
+        super().__init__()
+        self.stride, self.in_chs, self.out_chs = stride, in_chs, out_chs
+        mid_chs = make_divisible(out_chs * expand_ratio)
+        prenorm_act_layer = partial(get_norm_act_layer(norm_layer, act_layer), eps=norm_eps)
+
+        if stride == 2:
+            self.shortcut = Upsample2d(in_chs, out_chs)
+        elif in_chs != out_chs:
+            self.shortcut = nn.Conv2d(in_chs, out_chs, 1, bias=True)
+        else:
+            self.shortcut = nn.Identity()
+
+        self.pre_norm = prenorm_act_layer(in_chs, apply_act=False)
+
+        self.conv1_1x1 = create_conv2d(in_chs, mid_chs, 1, stride=1, bias=True)
+        self.act1 = _create_act(act_layer, inplace=True)
+        self.act2 = _create_act(act_layer, inplace=True)
+
+        self.up = Upsample2d(mid_chs, mid_chs) if stride == 2 else nn.Identity()
+        self.conv2_kxk = create_conv2d(mid_chs, mid_chs, kernel_size, stride=1, dilation=1, groups=mid_chs, bias=True)
+        self.conv3_1x1 = create_conv2d(mid_chs, out_chs, 1, bias=True)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def init_weights(self, scheme=''):
+        named_apply(partial(_init_conv, scheme=scheme), self)
+
+    def forward(self, x):
+        shortcut = self.shortcut(x)
+        x = self.pre_norm(x)
+
+        # 1x1 expansion conv & act
+        x = self.conv1_1x1(x)
+        x = self.act1(x)
+        x = self.up(x)
+
+        # (strided) depthwise 3x3 conv & act
+        x = self.conv2_kxk(x)
+        x = self.act2(x)
+
+        # 1x1 linear projection to output width
+        x = self.conv3_1x1(x)
+        x = self.drop_path(x) + shortcut
+
+        return x
+
+
+class InvStem(nn.Module):
+    def __init__(
+        self,
+        in_chs: int,
+        out_chs: int,
+        act_layer: str = 'gelu',
+        norm_layer: str = 'layernorm2d',
+        norm_eps: float = 1e-6,
+        bias: bool = True,
+    ):
+        super().__init__()
+        self.grad_checkpointing=False
+        norm_act_layer = partial(get_norm_act_layer(norm_layer, act_layer), eps=norm_eps)
+        self.out_chs = out_chs
+        self.conv1 = Upsample2d(in_chs, in_chs)
+        self.norm1 = norm_act_layer(in_chs)
+        self.conv2 = create_conv2d(in_chs, out_chs, 3, stride=1, bias=bias)
+        named_apply(_init_conv, self)
+
+    def forward(self, x):
+        if self.grad_checkpointing:
+            x = checkpoint(self.conv1, x)
+            x = self.norm1(x)
+            x = checkpoint(self.conv2, x)
+        else:
+            x = self.conv1(x)
+            x = self.norm1(x)
+            x = self.conv2(x)
+
+        return x
+
+
+class ViTaminDecoder(nn.Module):
+    def __init__(
+        self,
+        model,
+        num_query=0,
+        img_size=256,
+        drop_path=0.,
+        depths=(4, 2),
+        grad_ckpt=False,
+    ):
+        super().__init__()
+
+        self.num_query = num_query
+        vit = timm.create_model(
+            model,
+            fc_norm=False,
+            patch_size=1,
+            drop_rate=0.0,
+            num_classes=0,
+            global_pool='',
+            pos_embed='none',
+            mlp_layer=GeGluMlp,
+            class_token=False,
+            reg_tokens=num_query,
+            img_size=img_size,
+            drop_path_rate=drop_path,
+        )
+        self.blocks = vit.blocks
+        self.norm_pre = vit.norm_pre
+        self.norm = vit.norm
+
+        embed_dims = {
+            'vitamin_base': (768, 256, 128),
+            'vitamin_large': (1024, 320, 160)
+        }[model]
+        self.up_conv1 = Upsample2d(embed_dims[0], embed_dims[1])
+        self.up_conv2 = nn.Sequential(*[
+            InvMbConvLNBlock(
+                in_chs=embed_dims[1],
+                out_chs=embed_dims[1],
+                stride=2 if d == 0 else 1)
+            for d in range(depths[0])]
+        )
+        self.up_conv3 = nn.Sequential(*[
+            InvMbConvLNBlock(
+                in_chs=embed_dims[1] if d == 0 else embed_dims[2],
+                out_chs=embed_dims[2],
+                stride=2 if d == 0 else 1)
+            for d in range(depths[1])]
+        )
+        self.up_conv4 = InvStem(in_chs=embed_dims[2], out_chs=3)
+
+        self.grad_ckpt = grad_ckpt
+
+    def get_last_param(self):
+        return self.up_conv4.conv2.weight
+
+    def forward(self, x):
+        B, L, C = x.shape
+        H = W = int((L-self.num_query) ** 0.5)
+        x = self.norm_pre(x)
+        if self.grad_ckpt:
+            x = checkpoint_seq(self.blocks, x)
+            x = x[:, self.num_query:, :]
+            x = self.norm(x)
+            x = x.view(B, H, W, C).permute(0, 3, 1, 2)
+            x = checkpoint(self.up_conv1, x)
+            x = checkpoint_seq(self.up_conv2, x)
+            x = checkpoint_seq(self.up_conv3, x)
+        else:
+            x = self.blocks(x)
+            x = x[:, self.num_query:, :]
+            x = self.norm(x)
+            x = x.view(B, H, W, C).permute(0, 3, 1, 2)
+            x = self.up_conv1(x)
+            x = self.up_conv2(x)
+            x = self.up_conv3(x)
+        x = self.up_conv4(x)
+        return x
+
+
+def _create_vision_transformer(variant, pretrained=False, grad_ckpt=False, **kwargs) -> VisionTransformer:
+    if kwargs.get('features_only', None):
+        raise RuntimeError('features_only not implemented for Vision Transformer models.')
+
+    if 'flexi' in variant:
+        # FIXME Google FlexiViT pretrained models have a strong preference for bilinear patch / embed
+        # interpolation, other pretrained models resize better w/ anti-aliased bicubic interpolation.
+        _filter_fn = partial(checkpoint_filter_fn, interpolation='bilinear', antialias=False)
+    else:
+        _filter_fn = checkpoint_filter_fn
+
+    return build_model_with_cfg(
+        VisionTransformer,
+        variant,
+        pretrained,
+        pretrained_filter_fn=_filter_fn,
+        **kwargs,
+    )
+
+
+def _create_vision_transformer_hybrid(variant, backbone, pretrained=False, **kwargs):
+    embed_layer = partial(HybridEmbed, backbone=backbone)
+    kwargs.setdefault('patch_size', 1)  # default patch size for hybrid models if not set
+    return _create_vision_transformer(variant, pretrained=pretrained, embed_layer=embed_layer, **kwargs)
+
+
+@register_model
+def vitamin_small(pretrained=False, **kwargs) -> VisionTransformer:
+    stage_1_2 = MbConvStages(cfg=VitCfg(
+            embed_dim=(64, 128, 384),
+            depths=(2, 4, 1),
+            stem_width=64,
+            conv_cfg = VitConvCfg(
+                norm_layer='layernorm2d',
+                norm_eps=1e-6,
+            ),
+            head_type='1d',
+        ),
+    )
+    stage3_args = dict(embed_dim=384, depth=14, num_heads=6, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid('vitamin_small', backbone=stage_1_2, pretrained=pretrained, **dict(stage3_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_base(pretrained=False, **kwargs) -> VisionTransformer:
+    stage_1_2 = MbConvStages(cfg=VitCfg(
+            embed_dim=(128, 256, 768),
+            depths=(2, 4, 1),
+            stem_width=128,
+            conv_cfg = VitConvCfg(
+                norm_layer='layernorm2d',
+                norm_eps=1e-6,
+            ),
+            head_type='1d',
+        ),
+    )
+    stage3_args = dict(embed_dim=768, depth=14, num_heads=12, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid('vitamin_base', backbone=stage_1_2, pretrained=pretrained, **dict(stage3_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_base_256(pretrained=False, **kwargs) -> VisionTransformer:
+    stage_1_2 = MbConvStages(cfg=VitCfg(
+            embed_dim=(128, 256, 768),
+            depths=(2, 4, 1),
+            stem_width=128,
+            conv_cfg = VitConvCfg(
+                norm_layer='layernorm2d',
+                norm_eps=1e-6,
+            ),
+            head_type='1d',
+        ),
+    )
+    stage3_args = dict(img_size=256, embed_dim=768, depth=14, num_heads=12, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid('vitamin_base_256', backbone=stage_1_2, pretrained=pretrained, **dict(stage3_args, **kwargs))
+    return model
+
+
+@register_model
+def vitamin_large(pretrained=False, **kwargs) -> VisionTransformer:
+    stage_1_2 = MbConvStages(cfg=VitCfg(
+        embed_dim=(160, 320, 1024),
+        depths=(2, 4, 1),
+        stem_width=160,
+        conv_cfg = VitConvCfg(
+            norm_layer='layernorm2d',
+            norm_eps=1e-6,
+        ),
+        head_type='1d',
+    ),
+    )
+    stage3_args = dict(embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+    model = _create_vision_transformer_hybrid(
+        'vitamin_large', backbone=stage_1_2, pretrained=pretrained, **dict(stage3_args, **kwargs))
+    return model
+
+# @register_model
+# def vitamin_large_256(pretrained=False, **kwargs) -> VisionTransformer:
+#     backbone = MbConvStages(cfg=VitCfg(
+#         embed_dim=(160, 320, 1024),
+#         depths=(2, 4, 1),
+#         stem_width=160,
+#         conv_cfg = VitConvCfg(
+#             norm_layer='layernorm2d',
+#             norm_eps=1e-6,
+#         ),
+#         head_type='1d',
+#     ),
+#     )
+#     model_args = dict(img_size=256, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+#     model = _create_vision_transformer_hybrid(
+#         'vitamin_large_256', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+#     return model
+
+# @register_model
+# def vitamin_large_336(pretrained=False, **kwargs) -> VisionTransformer:
+#     backbone = MbConvStages(cfg=VitCfg(
+#         embed_dim=(160, 320, 1024),
+#         depths=(2, 4, 1),
+#         stem_width=160,
+#         conv_cfg = VitConvCfg(
+#             norm_layer='layernorm2d',
+#             norm_eps=1e-6,
+#         ),
+#         head_type='1d',
+#     ),
+#     )
+#     model_args = dict(img_size=336, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+#     model = _create_vision_transformer_hybrid(
+#         'vitamin_large_336', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+#     return model
+
+# @register_model
+# def vitamin_large_384(pretrained=False, **kwargs) -> VisionTransformer:
+#     backbone = MbConvStages(cfg=VitCfg(
+#         embed_dim=(160, 320, 1024),
+#         depths=(2, 4, 1),
+#         stem_width=160,
+#         conv_cfg = VitConvCfg(
+#             norm_layer='layernorm2d',
+#             norm_eps=1e-6,
+#         ),
+#         head_type='1d',
+#     ),
+#     )
+#     model_args = dict(img_size=384, embed_dim=1024, depth=31, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+#     model = _create_vision_transformer_hybrid(
+#         'vitamin_large_384', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+#     return model
+
+# @register_model
+# def vitamin_xlarge_256(pretrained=False, **kwargs) -> VisionTransformer:
+#     backbone = MbConvStages(cfg=VitCfg(
+#         embed_dim=(192, 384, 1152),
+#         depths=(2, 4, 1),
+#         stem_width=192,
+#         conv_cfg = VitConvCfg(
+#             norm_layer='layernorm2d',
+#             norm_eps=1e-6,
+#         ),
+#         head_type='1d',
+#     ),
+#     )
+#     model_args = dict(img_size=256, embed_dim=1152, depth=32, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+#     model = _create_vision_transformer_hybrid(
+#         'vitamin_xlarge_256', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+#     return model
+
+# @register_model
+# def vitamin_xlarge_336(pretrained=False, **kwargs) -> VisionTransformer:
+#     backbone = MbConvStages(cfg=VitCfg(
+#         embed_dim=(192, 384, 1152),
+#         depths=(2, 4, 1),
+#         stem_width=192,
+#         conv_cfg = VitConvCfg(
+#             norm_layer='layernorm2d',
+#             norm_eps=1e-6,
+#         ),
+#         head_type='1d',
+#     ),
+#     )
+#     model_args = dict(img_size=336, embed_dim=1152, depth=32, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+#     model = _create_vision_transformer_hybrid(
+#         'vitamin_xlarge_256', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+#     return model
+
+# @register_model
+# def vitamin_xlarge_384(pretrained=False, **kwargs) -> VisionTransformer:
+#     backbone = MbConvStages(cfg=VitCfg(
+#         embed_dim=(192, 384, 1152),
+#         depths=(2, 4, 1),
+#         stem_width=192,
+#         conv_cfg = VitConvCfg(
+#             norm_layer='layernorm2d',
+#             norm_eps=1e-6,
+#         ),
+#         head_type='1d',
+#     ),
+#     )
+#     model_args = dict(img_size=384, embed_dim=1152, depth=32, num_heads=16, mlp_layer=GeGluMlp, mlp_ratio=2., class_token=False, global_pool='avg')
+#     model = _create_vision_transformer_hybrid(
+#         'vitamin_xlarge_384', backbone=backbone, pretrained=pretrained, **dict(model_args, **kwargs))
+#     return model
+
+
+def count_params(model: nn.Module):
+    return sum([m.numel() for m in model.parameters()])
+
+
+def count_stage_params(model: nn.Module, prefix='none'):
+    collections = []
+    for name, m in model.named_parameters():
+        print(name)
+        if name.startswith(prefix):
+            collections.append(m.numel())
+    return sum(collections)
+
+
+if __name__ == "__main__":
+    # ViTaminDecoder('vitamin_base', img_size=256, patch_size=16)
+    # model = timm.create_model(
+    #     'vitamin_base',
+    #     fc_norm=True,
+    #     drop_rate=0.0,
+    #     num_classes=0,
+    #     global_pool='',
+    #     mlp_layer=GeGluMlp,
+    #     class_token=False,
+    #     reg_tokens=32,
+    #     img_size=256,
+    #     patch_size=1,
+    #     drop_path_rate=0.1,
+    # )
+    # print(model.has_class_token)
+    # print(model.num_prefix_tokens)
+    # print(model.pos_embed.shape)
+    Stem(64, 64)

native/models/vqvae.py

@@ -0,0 +1,175 @@
+import timm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from contextlib import nullcontext
+from torch.nn.functional import scaled_dot_product_attention
+
+from native.models.quant import VectorQuantizerM
+from native.models.vitamin import ViTaminDecoder, GeGluMlp
+
+
+class PlainAttention(nn.Module):
+    def __init__(self, in_dim, out_dim, num_heads):
+        super().__init__()
+        if in_dim > out_dim:
+            # assert in_dim // num_heads == out_dim
+            self.head_dim = in_dim // num_heads
+            self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False)
+            self.q_bias = nn.Parameter(torch.zeros(in_dim))
+            self.v_bias = nn.Parameter(torch.zeros(in_dim))
+            self.register_buffer('zero_k_bias', torch.zeros(in_dim))
+        else:
+            # assert out_dim // num_heads == in_dim
+            self.head_dim = out_dim // num_heads
+            self.qkv = nn.Linear(in_dim, out_dim * 3, bias=False)
+            self.q_bias = nn.Parameter(torch.zeros(out_dim))
+            self.v_bias = nn.Parameter(torch.zeros(out_dim))
+            self.register_buffer('zero_k_bias', torch.zeros(out_dim))
+
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.scale = self.head_dim ** -0.5
+        self.proj = nn.Linear(out_dim, out_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)))
+        q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+
+        x = scaled_dot_product_attention(q, k, v)
+
+        if self.in_dim > self.out_dim:
+            x = torch.mean(x, dim=1)
+            if self.in_dim // self.num_heads != self.out_dim:
+                x = nn.functional.adaptive_avg_pool1d(x, self.out_dim)
+        else:
+            x = x.transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        return x
+
+
+class AttnProjection(nn.Module):
+    def __init__(self, in_dim, out_dim, num_heads, norm_layer=nn.LayerNorm, mlp_ratio=2):
+        super().__init__()
+        assert out_dim % in_dim == 0 or in_dim % out_dim == 0
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.norm1 = norm_layer(in_dim)
+        self.attn = PlainAttention(in_dim, out_dim, num_heads)
+        self.proj = nn.Linear(in_dim, out_dim)
+        self.norm3 = norm_layer(in_dim)
+
+        self.norm2 = norm_layer(out_dim)
+        hidden_dim = int(out_dim * mlp_ratio)
+        self.mlp = GeGluMlp(
+            in_features=out_dim,
+            hidden_features=hidden_dim
+        )
+
+    def forward(self, x):
+        x = self.proj(self.norm3(x)) + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class VQVAE(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        # 1. build encoder
+        self.encoder = timm.create_model(
+            args.model,
+            patch_size=1,
+            fc_norm=True,
+            drop_rate=0.0,
+            num_classes=0,
+            global_pool='',
+            pos_embed='none',
+            class_token=False,
+            mlp_layer=GeGluMlp,
+            img_size=args.img_size,
+            drop_path_rate=args.drop_path,
+        )
+        self.encoder.set_grad_checkpointing(args.grad_ckpt)
+
+        # 2. build conv before quant
+        if args.quant_proj == 'linear':
+            self.quant_proj = nn.Linear(self.encoder.embed_dim, args.vocab_width)
+        elif args.quant_proj == 'attn':
+            self.quant_proj = AttnProjection(self.encoder.embed_dim, args.vocab_width, args.num_codebooks)
+        else:
+            raise NotImplementedError
+
+        # 3. build quant
+        self.quantize = VectorQuantizerM(
+            vocab_size=args.vocab_size,
+            vocab_width=args.vocab_width,
+            beta=args.vq_beta,
+            use_entropy_loss=args.le > 0,
+            entropy_temp=args.e_temp,
+            num_codebooks=args.num_codebooks,
+        )
+
+        # 4. build conv after quant
+        if args.quant_proj == 'linear':
+            self.post_quant_proj = nn.Linear(args.vocab_width, self.encoder.embed_dim)
+        elif args.quant_proj == 'attn':
+            self.post_quant_proj = AttnProjection(args.vocab_width, self.encoder.embed_dim, args.num_codebooks)
+        else:
+            raise NotImplementedError
+
+        # 5. build decoder
+        self.decoder = ViTaminDecoder(
+            args.model,
+            depths=(4, 2),
+            img_size=args.img_size,
+            drop_path=args.drop_path,
+            grad_ckpt=args.grad_ckpt
+        )
+
+        self.maybe_record_function = nullcontext
+
+    def forward(self, img):
+        features = self.encoder(img).float()
+        with torch.cuda.amp.autocast(enabled=False):
+            features = self.quant_proj(features)
+            quant_out = self.quantize(features)
+            features, vq_loss, entropy_loss, usages = quant_out
+            features = self.post_quant_proj(features)
+        rec_img = self.decoder(features).float()
+        return rec_img, vq_loss, entropy_loss, usages
+
+    def img_to_idx(self, img):
+        features = self.encoder(img).float()
+        features = self.quant_proj(features)
+        return self.quantize.f_to_idx(features)
+
+    def idx_to_img(self, indices):
+        features = self.quantize.idx_to_f(indices)
+        features = self.post_quant_proj(features)
+        img = self.decoder(features).clamp_(-1, 1)
+        return img
+
+    def img_to_reconstructed_img(self, img) -> torch.Tensor:
+        features = self.encoder(img).float()
+        with torch.cuda.amp.autocast(enabled=False):
+            features = self.quant_proj(features)
+            quant_out = self.quantize(features)
+            features, _, _, _ = quant_out
+            features = self.post_quant_proj(features)
+        rec_img = self.decoder(features).float().clamp_(-1, 1)
+        return rec_img
+
+
+if __name__ == '__main__':
+    for clz in (nn.Linear, nn.LayerNorm, nn.BatchNorm2d, nn.SyncBatchNorm, nn.Conv1d, nn.Conv2d, nn.ConvTranspose1d,
+                nn.ConvTranspose2d):
+        setattr(clz, 'reset_parameters', lambda self: None)
+
+    cnn = VQVAE(channel_num=64, vocab_norm=False)
+    from models import init_weights
+
+    init_weights(cnn, -0.5)
+    torch.save(cnn.state_dict(), r'C:\Users\16333\Desktop\PyCharm\vlip\local_output\cnn.pth')

native/open_clip/source_model.py

@@ -0,0 +1,599 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+from dataclasses import dataclass
+import logging
+import math
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from .hf_model import HFTextEncoder
+from .modified_resnet import ModifiedResNet
+from .transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+from .utils import to_2tuple
+from .timm_model import TimmModel 
+
+
+import time
+
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    input_patchnorm: bool = False  # whether to use dual patchnorm - would only apply the input layernorm on each patch, as post-layernorm already exist in original clip vit design
+    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer
+    n_queries: int = 256  # n_queries for attentional pooler
+    attn_pooler_heads: int = 8  # n heads for attentional_pooling
+    output_tokens: bool = False
+
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+    embed_cls: bool = False
+    pad_id: int = 0
+    output_tokens: bool = False
+    text_mask: str = 'first' # default first truncate in bpe_tokenizer
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            patch_drop=vision_cfg.patch_dropout if vision_cfg.patch_dropout > 0 else None,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            input_patchnorm=vision_cfg.input_patchnorm,
+            global_average_pool=vision_cfg.global_average_pool,
+            attentional_pool=vision_cfg.attentional_pool,
+            n_queries=vision_cfg.n_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj=text_cfg.proj,
+            pooler_type=text_cfg.pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            output_tokens=text_cfg.output_tokens,
+            pad_id=text_cfg.pad_id,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.context_length = text.context_length
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.method_lock_text_tower = text.lock
+        self.text_no_grad = False
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True, unlock_text_proj=False):
+        # added by jieneng
+        self.method_lock_text_tower(unlocked_layers, freeze_layer_norm)
+        self.text_no_grad = True
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True, enable_text=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable_text
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual.trunk.forward_features(image)
+        features = self.visual.trunk.forward_head(features)
+        features = self.visual.head(features)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        torch.cuda.synchronize()
+        # time_s1 = time.time()
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        # torch.cuda.synchronize()
+        # time_e1 = time.time()
+        # print("image ", time_e1-time_s1)
+
+        if self.text_no_grad:
+            with torch.no_grad():
+                text_features = self.encode_text(text, normalize=True).detach() if text is not None else None
+        else:
+            text_features = self.encode_text(text, normalize=True) if text is not None else None
+
+        # torch.cuda.synchronize()
+        # time_e2 = time.time()
+        # print("text ", time_e2-time_e1)
+
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomTextCLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.context_length = self.text.context_length
+        self.vocab_size = self.text.vocab_size
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        self.text_no_grad = False
+        # self.context_length = 0 in t5
+        # assert 1==2, (self.text, self.context_length, self.vocab_size)
+
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True, unlock_text_proj = False):
+        self.text.lock(unlocked_layers, freeze_layer_norm, unlock_text_proj)
+        self.text_no_grad = True
+
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True, enable_text=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable_text)
+
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual.trunk.forward_features(image)
+        features = self.visual.trunk.forward_head(features)
+        features = self.visual.head(features)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        # if self.text_no_grad:
+        #     with torch.no_grad():
+        #         text_features = self.encode_text(text, normalize=True).detach() if text is not None else None
+        # else:
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, (CLIP, TextTransformer)):
+            # convert text nn.Parameter projections
+            attr = getattr(l, "text_projection", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+        if isinstance(l, VisionTransformer):
+            # convert vision nn.Parameter projections
+            attr = getattr(l, "proj", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+def resize_pos_embed_timm(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.trunk.pos_embed', None) # 1, 196, 1024]
+    if old_pos_embed is None:
+        return
+    
+    grid_size = to_2tuple(model.visual.trunk.patch_embed.grid_size)
+ 
+    
+    if hasattr(model.visual.trunk, 'cls_token') and model.visual.trunk.cls_token is not None:
+        return
+        # extra_tokens?
+        raise NotImplementedError
+
+    new_seq_len = grid_size[0] * grid_size[1]
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    pos_emb_img = old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img[0]))))
+    old_pos_emb_img = pos_emb_img
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size) # Resizing position embedding grid-size from (1, 1) to (21, 21)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)
+    state_dict['visual.trunk.pos_embed'] = pos_emb_img
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    pe_key_name = 'visual.positional_embedding'
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None:
+        pe_key_name = 'visual.trunk.pos_embed'
+        old_pos_embed = state_dict.get('visual.trunk.pos_embed', None) # 1, 196, 1024]
+        
+    if old_pos_embed is None:
+        return
+    
+    if hasattr(model.visual, 'grid_size'):
+        grid_size = to_2tuple(model.visual.grid_size)
+    elif hasattr(model.visual.trunk.patch_embed, 'grid_size'):
+        grid_size = to_2tuple(model.visual.trunk.patch_embed.grid_size)
+    else:
+        return
+    
+    if hasattr(model.visual.trunk, 'cls_token') and model.visual.trunk.cls_token is not None:
+        extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    else:
+        extra_tokens = 0
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+    old_pos_emb_img = pos_emb_img
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size) # Resizing position embedding grid-size from (1, 1) to (21, 21)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+
+
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict[pe_key_name] = new_pos_embed
+
+def resize_text_pos_embed(state_dict, model, interpolation: str = 'linear', antialias: bool = False):
+    old_pos_embed = state_dict.get('positional_embedding', None)
+    if old_pos_embed is None:
+        return
+    # FIXME add support for text cls_token
+    model_pos_embed = getattr(model, 'positional_embedding', None)
+    if model_pos_embed is None:
+        model_pos_embed = getattr(model.text, 'positional_embedding', None)
+
+    old_num_pos = old_pos_embed.shape[0]
+    old_width = old_pos_embed.shape[1]
+    num_pos = model_pos_embed.shape[0]
+    width = model_pos_embed.shape[1]
+    assert old_width == width, 'text pos_embed width changed!'
+    if old_num_pos == num_pos:
+        return
+
+    logging.info('Resizing text position embedding num_pos from %s to %s', old_num_pos, num_pos)
+    old_pos_embed = old_pos_embed.reshape(1, old_num_pos, old_width).permute(0, 2, 1)
+    old_pos_embed = F.interpolate(
+        old_pos_embed,
+        size=num_pos,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    old_pos_embed = old_pos_embed.permute(0, 2, 1)[0]
+    new_pos_embed = old_pos_embed
+
+    state_dict['positional_embedding'] = new_pos_embed

native/open_clip/transformer.py

@@ -0,0 +1,735 @@
+from collections import OrderedDict
+import math
+from typing import Callable, Optional, Sequence, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+from .utils import to_2tuple
+
+
+class LayerNormFp32(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x.to(torch.float32), self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm (with cast back to input dtype)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class QuickGELU(nn.Module):
+    # NOTE This is slower than nn.GELU or nn.SiLU and uses more GPU memory
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim,
+            num_heads=8,
+            qkv_bias=True,
+            scaled_cosine=False,
+            scale_heads=False,
+            logit_scale_max=math.log(1. / 0.01),
+            attn_drop=0.,
+            proj_drop=0.
+    ):
+        super().__init__()
+        self.scaled_cosine = scaled_cosine
+        self.scale_heads = scale_heads
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.logit_scale_max = logit_scale_max
+
+        # keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
+        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
+        if qkv_bias:
+            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
+        else:
+            self.in_proj_bias = None
+
+        if self.scaled_cosine:
+            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
+        else:
+            self.logit_scale = None
+        self.attn_drop = nn.Dropout(attn_drop)
+        if self.scale_heads:
+            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
+        else:
+            self.head_scale = None
+        self.out_proj = nn.Linear(dim, dim)
+        self.out_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
+        L, N, C = x.shape
+        q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
+        q = q.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+        k = k.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+        v = v.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+
+        if self.logit_scale is not None:
+            attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
+            logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
+            attn = attn.view(N, self.num_heads, L, L) * logit_scale
+            attn = attn.view(-1, L, L)
+        else:
+            q = q * self.scale
+            attn = torch.bmm(q, k.transpose(-1, -2))
+
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+                new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+                attn_mask = new_attn_mask
+            attn += attn_mask
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = torch.bmm(attn, v)
+        if self.head_scale is not None:
+            x = x.view(N, self.num_heads, L, C) * self.head_scale
+            x = x.view(-1, L, C)
+        x = x.transpose(0, 1).reshape(L, N, C)
+        x = self.out_proj(x)
+        x = self.out_drop(x)
+        return x
+
+
+class AttentionalPooler(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            context_dim: int,
+            n_head: int = 8,
+            n_queries: int = 256,
+            norm_layer: Callable = LayerNorm
+    ):
+        super().__init__()
+        self.query = nn.Parameter(torch.randn(n_queries, d_model))
+        self.attn = nn.MultiheadAttention(d_model, n_head, kdim=context_dim, vdim=context_dim)
+        self.ln_q = norm_layer(d_model)
+        self.ln_k = norm_layer(context_dim)
+
+    def forward(self, x: torch.Tensor):
+        x = self.ln_k(x).permute(1, 0, 2)  # NLD -> LND
+        N = x.shape[1]
+        q = self.ln_q(self.query)
+        out = self.attn(self._repeat(q, N), x, x, need_weights=False)[0]
+        return out.permute(1, 0, 2)  # LND -> NLD
+
+    def _repeat(self, query, N: int):
+        return query.unsqueeze(1).repeat(1, N, 1)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            is_cross_attention: bool = False,
+    ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+        if is_cross_attention:
+            self.ln_1_kv = norm_layer(d_model)
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ("gelu", act_layer()),
+            ("c_proj", nn.Linear(mlp_width, d_model))
+        ]))
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+    def attention(
+            self,
+            q_x: torch.Tensor,
+            k_x: Optional[torch.Tensor] = None,
+            v_x: Optional[torch.Tensor] = None,
+            attn_mask: Optional[torch.Tensor] = None,
+    ):
+        k_x = k_x if k_x is not None else q_x
+        v_x = v_x if v_x is not None else q_x
+
+        attn_mask = attn_mask.to(q_x.dtype) if attn_mask is not None else None
+        return self.attn(
+            q_x, k_x, v_x, need_weights=False, attn_mask=attn_mask
+        )[0]
+
+    def forward(
+            self,
+            q_x: torch.Tensor,
+            k_x: Optional[torch.Tensor] = None,
+            v_x: Optional[torch.Tensor] = None,
+            attn_mask: Optional[torch.Tensor] = None,
+    ):
+        k_x = self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None
+        v_x = self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None
+
+        x = q_x + self.ls_1(self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask))
+        x = x + self.ls_2(self.mlp(self.ln_2(x)))
+        return x
+
+
+class CustomResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            scale_cosine_attn: bool = False,
+            scale_heads: bool = False,
+            scale_attn: bool = False,
+            scale_fc: bool = False,
+    ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        self.attn = Attention(
+            d_model, n_head,
+            scaled_cosine=scale_cosine_attn,
+            scale_heads=scale_heads,
+        )
+        self.ln_attn = norm_layer(d_model) if scale_attn else nn.Identity()
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ('ln', norm_layer(mlp_width) if scale_fc else nn.Identity()),
+            ("gelu", act_layer()),
+            ("c_proj", nn.Linear(mlp_width, d_model))
+        ]))
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        x = x + self.ls_1(self.ln_attn(self.attn(self.ln_1(x), attn_mask=attn_mask)))
+        x = x + self.ls_2(self.mlp(self.ln_2(x)))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.grad_checkpointing = False
+
+        self.resblocks = nn.ModuleList([
+            ResidualAttentionBlock(
+                width, heads, mlp_ratio, ls_init_value=ls_init_value, act_layer=act_layer, norm_layer=norm_layer)
+            for _ in range(layers)
+        ])
+
+    def get_cast_dtype(self) -> torch.dtype:
+        if hasattr(self.resblocks[0].mlp.c_fc, 'int8_original_dtype'):
+            return self.resblocks[0].mlp.c_fc.int8_original_dtype
+        return self.resblocks[0].mlp.c_fc.weight.dtype
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        for r in self.resblocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                # TODO: handle kwargs https://github.com/pytorch/pytorch/issues/79887#issuecomment-1161758372
+                x = checkpoint(r, x, None, None, attn_mask)
+            else:
+                x = r(x, attn_mask=attn_mask)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            image_size: int,
+            patch_size: int,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float,
+            ls_init_value: float = None,
+            global_average_pool: bool = False,
+            attentional_pool: bool = False,
+            n_queries: int = 256,
+            attn_pooler_heads: int = 8,
+            output_dim: int = 512,
+            patch_dropout: float = 0.,
+            input_patchnorm: bool = False,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            output_tokens: bool = False
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        image_height, image_width = self.image_size = to_2tuple(image_size)
+        patch_height, patch_width = self.patch_size = to_2tuple(patch_size)
+        self.grid_size = (image_height // patch_height, image_width // patch_width)
+        self.output_dim = output_dim
+
+        # whether to layernorm each patch, as done in dual patchnorm paper - https://arxiv.org/abs/2302.01327v1
+        self.input_patchnorm = input_patchnorm
+
+        if input_patchnorm:
+            patch_input_dim = patch_height * patch_width * 3
+            self.patchnorm_pre_ln = LayerNorm(patch_input_dim)
+            self.conv1 = nn.Linear(patch_input_dim, width)
+        else:
+            self.patchnorm_pre_ln = nn.Identity()
+            self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size,
+                                   bias=False)
+
+        # class embeddings and positional embeddings
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn(self.grid_size[0] * self.grid_size[1] + 1, width))
+
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+        self.ln_pre = norm_layer(width)
+        self.transformer = Transformer(
+            width,
+            layers,
+            heads,
+            mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+        self.global_average_pool = global_average_pool
+        if attentional_pool:
+            self.attn_pool = AttentionalPooler(output_dim, width, n_head=attn_pooler_heads, n_queries=n_queries)
+            self.ln_post = norm_layer(output_dim)
+            self.proj = nn.Parameter(scale * torch.randn(output_dim, output_dim))
+        else:
+            self.attn_pool = None
+            self.ln_post = norm_layer(width)
+            self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+        self.init_parameters()
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        for param in self.parameters():
+            param.requires_grad = False
+
+        if unlocked_groups != 0:
+            groups = [
+                [
+                    self.conv1,
+                    self.class_embedding,
+                    self.positional_embedding,
+                    self.ln_pre,
+                ],
+                *self.transformer.resblocks[:-1],
+                [
+                    self.transformer.resblocks[-1],
+                    self.ln_post,
+                ],
+                self.proj,
+            ]
+
+            def _unlock(x):
+                if isinstance(x, Sequence):
+                    for g in x:
+                        _unlock(g)
+                else:
+                    if isinstance(x, torch.nn.Parameter):
+                        x.requires_grad = True
+                    else:
+                        for p in x.parameters():
+                            p.requires_grad = True
+
+            _unlock(groups[-unlocked_groups:])
+
+    def init_parameters(self):
+        # FIXME OpenAI CLIP did not define an init for the VisualTransformer
+        # TODO experiment if default PyTorch init, below, or alternate init is best.
+
+        # nn.init.normal_(self.class_embedding, std=self.scale)
+        # nn.init.normal_(self.positional_embedding, std=self.scale)
+        #
+        # proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        # attn_std = self.transformer.width ** -0.5
+        # fc_std = (2 * self.transformer.width) ** -0.5
+        # for block in self.transformer.resblocks:
+        #     nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+        #     nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+        #     nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+        #     nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+        #
+        # if self.text_projection is not None:
+        #     nn.init.normal_(self.text_projection, std=self.scale)
+        pass
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.grad_checkpointing = enable
+
+    def _global_pool(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.global_average_pool:
+            return x.mean(dim=1), x
+        else:
+            return x[:, 0], x[:, 1:]
+
+    def forward(self, x: torch.Tensor):
+
+        # to patches - whether to use dual patchnorm - https://arxiv.org/abs/2302.01327v1
+        if self.input_patchnorm:
+            # einops - rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)')
+            x = x.reshape(x.shape[0], x.shape[1], self.grid_size[0], self.patch_size[0], self.grid_size[1],
+                          self.patch_size[1])
+            x = x.permute(0, 2, 4, 1, 3, 5)
+            x = x.reshape(x.shape[0], self.grid_size[0] * self.grid_size[1], -1)
+            x = self.patchnorm_pre_ln(x)
+            x = self.conv1(x)
+        else:
+            x = self.conv1(x)  # shape = [*, width, grid, grid]
+            x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+            x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+
+        # class embeddings and positional embeddings
+        x = torch.cat(
+            [self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
+             x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        x = self.patch_dropout(x)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        if self.attn_pool is not None:
+            x = self.attn_pool(x)
+            x = self.ln_post(x)
+            pooled, tokens = self._global_pool(x)
+        else:
+            pooled, tokens = self._global_pool(x)
+            pooled = self.ln_post(pooled)
+
+        if self.proj is not None:
+            pooled = pooled @ self.proj
+
+        if self.output_tokens:
+            return pooled, tokens
+
+        return pooled
+
+
+class TextTransformer(nn.Module):
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            context_length: int = 77,
+            vocab_size: int = 49408,
+            width: int = 512,
+            heads: int = 8,
+            layers: int = 12,
+            ls_init_value: float = None,
+            output_dim: int = 512,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            embed_cls: bool = False,
+            pad_id: int = 0,
+            output_tokens: bool = False,
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        self.num_pos = self.context_length = context_length
+        self.vocab_size = vocab_size
+        self.width = width
+        self.output_dim = output_dim
+        self.heads = heads
+        self.pad_id = pad_id
+
+        self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+
+        if embed_cls:
+            self.cls_emb = nn.Parameter(torch.empty(width))
+            self.num_pos += 1
+        else:
+            self.cls_emb = None
+
+        self.token_embedding = nn.Embedding(vocab_size, width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.num_pos, width))
+        self.transformer = Transformer(
+            width=width,
+            layers=layers,
+            heads=heads,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+        self.ln_final = norm_layer(width)
+
+        self.register_buffer('attn_mask', self.build_attention_mask(), persistent=False)
+
+        self.init_parameters()
+
+    def lock(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True, unlock_text_proj=False):
+        for param in self.parameters():
+            param.requires_grad = False
+
+        if unlocked_layers == -1:
+            self.text_projection.requires_grad = True
+
+    def init_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+        if self.cls_emb is not None:
+            nn.init.normal_(self.cls_emb, std=0.01)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.grad_checkpointing = enable
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.num_pos, self.num_pos)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    def build_cls_mask(self, text, cast_dtype: torch.dtype):
+        cls_mask = (text != self.pad_id).unsqueeze(1)
+        cls_mask = F.pad(cls_mask, (1, 0, cls_mask.shape[2], 0), value=1.0)
+        additive_mask = torch.empty(cls_mask.shape, dtype=cast_dtype, device=cls_mask.device)
+        additive_mask.fill_(0)
+        additive_mask.masked_fill_(~cls_mask, float("-inf"))
+        additive_mask = torch.repeat_interleave(additive_mask, self.heads, 0)
+        return additive_mask
+
+    def _repeat(self, t, N: int):
+        return t.reshape(1, 1, -1).repeat(N, 1, 1)
+
+    def forward(self, text):
+        cast_dtype = self.transformer.get_cast_dtype()
+        seq_len = text.shape[1]
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+        attn_mask = self.attn_mask
+        if self.cls_emb is not None:
+            seq_len += 1
+            x = torch.cat([x, self._repeat(self.cls_emb, x.shape[0])], dim=1)
+            cls_mask = self.build_cls_mask(text, cast_dtype)
+            attn_mask = attn_mask[None, :seq_len, :seq_len] + cls_mask[:, :seq_len, :seq_len]
+
+        x = x + self.positional_embedding[:seq_len].to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        if self.cls_emb is not None:
+            pooled, tokens = x[:, -1], x[:, :-1]
+            pooled = self.ln_final(pooled)
+        else:
+            x = self.ln_final(x)
+            pooled, tokens = x[torch.arange(x.shape[0]), text.argmax(dim=-1)], x
+
+        if self.text_projection is not None:
+            pooled = pooled @ self.text_projection
+
+        if self.output_tokens:
+            return pooled, tokens
+
+        return pooled
+
+
+class MultimodalTransformer(Transformer):
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            context_length: int = 77,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            output_dim: int = 512,
+    ):
+
+        super().__init__(
+            width=width,
+            layers=layers,
+            heads=heads,
+            mlp_ratio=mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+        self.context_length = context_length
+        self.cross_attn = nn.ModuleList([
+            ResidualAttentionBlock(
+                width,
+                heads,
+                mlp_ratio,
+                ls_init_value=ls_init_value,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                is_cross_attention=True,
+            )
+            for _ in range(layers)
+        ])
+
+        self.register_buffer('attn_mask', self.build_attention_mask(), persistent=False)
+
+        self.ln_final = norm_layer(width)
+        self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+
+    def init_parameters(self):
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+        for block in self.transformer.cross_attn:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    def forward(self, image_embs, text_embs):
+        text_embs = text_embs.permute(1, 0, 2)  # NLD -> LNDsq
+        image_embs = image_embs.permute(1, 0, 2)  # NLD -> LND
+        seq_len = text_embs.shape[0]
+
+        for resblock, cross_attn in zip(self.resblocks, self.cross_attn):
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                # TODO: handle kwargs https://github.com/pytorch/pytorch/issues/79887#issuecomment-1161758372
+                text_embs = checkpoint(resblock, text_embs, None, None, self.attn_mask[:seq_len, :seq_len])
+                text_embs = checkpoint(cross_attn, text_embs, image_embs, image_embs, None)
+            else:
+                text_embs = resblock(text_embs, attn_mask=self.attn_mask[:seq_len, :seq_len])
+                text_embs = cross_attn(text_embs, k_x=image_embs, v_x=image_embs)
+
+        x = text_embs.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)
+
+        if self.text_projection is not None:
+            x = x @ self.text_projection
+
+        return x
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable

native/open_clip/utils.py

@@ -0,0 +1,89 @@
+from itertools import repeat
+import collections.abc
+
+import torch
+from torch import nn as nn
+from torchvision.ops.misc import FrozenBatchNorm2d
+
+
+def freeze_batch_norm_2d(module, module_match={}, name=''):
+    """
+    Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is
+    itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and
+    returned. Otherwise, the module is walked recursively and submodules are converted in place.
+
+    Args:
+        module (torch.nn.Module): Any PyTorch module.
+        module_match (dict): Dictionary of full module names to freeze (all if empty)
+        name (str): Full module name (prefix)
+
+    Returns:
+        torch.nn.Module: Resulting module
+
+    Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762
+    """
+    res = module
+    is_match = True
+    if module_match:
+        is_match = name in module_match
+    if is_match and isinstance(module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)):
+        res = FrozenBatchNorm2d(module.num_features)
+        res.num_features = module.num_features
+        res.affine = module.affine
+        if module.affine:
+            res.weight.data = module.weight.data.clone().detach()
+            res.bias.data = module.bias.data.clone().detach()
+        res.running_mean.data = module.running_mean.data
+        res.running_var.data = module.running_var.data
+        res.eps = module.eps
+    else:
+        for child_name, child in module.named_children():
+            full_child_name = '.'.join([name, child_name]) if name else child_name
+            new_child = freeze_batch_norm_2d(child, module_match, full_child_name)
+            if new_child is not child:
+                res.add_module(child_name, new_child)
+    return res
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = lambda n, x: _ntuple(n)(x)
+
+# Replaces all linear layers with linear_replacement
+# TODO: add int8 support for other linear layers including attn and convnets
+def replace_linear(model, linear_replacement, include_modules=['c_fc', 'c_proj'], copy_weights=True):
+    for name, module in model.named_children():
+        if len(list(module.children())) > 0:
+            replace_linear(module, linear_replacement, include_modules, copy_weights)
+
+        if isinstance(module, torch.nn.Linear) and name in include_modules:
+            old_module = model._modules[name]
+            model._modules[name] = linear_replacement(
+                module.in_features,
+                module.out_features,
+                module.bias is not None,
+            )
+            if copy_weights:
+                model._modules[name].weight.data.copy_(old_module.weight.data)
+                if model._modules[name].bias is not None:
+                    model._modules[name].bias.data.copy_(old_module.bias)
+
+    return model
+
+def convert_int8_model_to_inference_mode(model):
+    for m in model.modules():
+        if hasattr(m, 'prepare_for_eval'):
+            int8_original_dtype = m.weight.dtype
+            m.prepare_for_eval()
+            m.int8_original_dtype = int8_original_dtype

native/utils/config.py

@@ -0,0 +1,243 @@
+import os
+import sys
+import torch
+import random
+import numpy as np
+from tap import Tap
+from typing import Optional, Union
+from collections import OrderedDict
+
+from native.utils import dist
+
+
+class Args(Tap):
+    model: str = 'vitamin_large' # 'vitamin_base', 'vitamin_large', xxx
+    exp_name: str = 'unitok_large'
+    output_dir: str = 'local_output'
+    resume_from: str = ''  # if specified, load this checkpoint; if not, load the latest checkpoint in output_dir (if exists)
+    lpips_path: str = 'external/lpips_with_vgg.pth'
+    dino_path: str = 'external/dinov2_vits14_pretrain.pth'
+    fid_eval_src: str = ''
+    fid_eval_dst: str = ''
+    vis_img_dir: str = 'asset/vis_imgs/'
+    fid_feature_extractor: str = 'external/weights-inception-2015-12-05-6726825d.pth'
+    clip_pretrain_path: str = ''
+
+    # speed-up
+    fp16: bool = False  # whether to use FP16
+    bf16: bool = True  # whether to use BF16
+    tf32: bool = True  # whether to use TensorFloat32
+    compile_model: bool = False  # whether to use torch.compile()
+    ddp_static: bool = False  # whether to use static graph in DDP
+    grad_ckpt: bool = True  # gradient checkpointing
+    grad_accu: int = 1  # gradient accumulation
+    device: str = 'cpu' # will be set automatically
+    dtype: torch.dtype = torch.float32 # will be set automatically
+
+    # data
+    train_data: str = None
+    val_data: str = None
+    dataset_type: str = 'webdataset'
+    imagenet_val: str = None
+    imagenet_v2: str = None
+    subset_ratio: float = 1.0
+    img_size: int = 256
+    resize_ratio: float = 1.125  # only applicable to 'img' dataset_type
+    hflip: bool = False
+    workers: int = 8  # num workers; 0: auto, -1: don't use multiprocessing in DataLoader
+    train_num_samples: int = 1280_000_000
+    train_data_upsampling_factors: str = None
+    dataset_resampled: bool = False
+    use_aug: bool = False
+
+    # quantizer
+    vocab_size: int = 32768
+    vocab_width: int = 64
+    vocab_norm: bool = True
+    vq_beta: float = 0.25  # commitment loss weight
+    num_codebooks: int = 8
+    quant_proj: str = 'attn'
+
+    # model
+    embed_dim: int = 768
+    num_query: int = 0
+    use_clip_pretrain: bool = False
+    patch_size: int = 16
+    drop_path: float = 0.1
+    text_width: int = 768
+    text_heads: int = 12
+    text_layers: int = 12
+    text_vocab_size: int = 49408
+    text_context_length: int = 77
+
+    # CLIP
+    local_loss: bool = True
+    gather_with_grad: bool = True
+    pretrained_clip: str = None
+    pretrained_clip_text: str = None
+    lock_text: bool = False
+    lock_text_unlocked_layers: int = 0
+    lock_text_freeze_layer_norm: bool = False
+    force_custom_text: bool = False
+    force_custom_vision: bool = False
+    zeroshot_eval_freq: int = 1
+
+    # discriminator
+    dino_depth: int = 12
+    dino_kernel_size: int = 9
+    disc_norm: str = 'gn'  # gn: group norm, bn: batch norm, sbn: sync batch norm, hbn: hybrid sync batch norm
+    disc_aug_prob: float = 1.0
+    disc_specnorm: bool = False
+    step_disc_every: int = 1
+
+    # initialization
+    vae_init: float = -0.5  # <0: xavier_normal_(gain=abs(init)); >0: trunc_normal_(std=init)
+    vocab_init: float = -1  # <0: uniform(-abs(init)*base, abs(init)*base), where base = 20/vocab_size; >0: trunc_normal_(std=init)
+    disc_init: float = -0.5  # <0: xavier_normal_(gain=abs(init)); >0: trunc_normal_(std=init)
+
+    # optimization
+    epoch: int = 1  # number of epochs
+    local_bs: int = 64  # batch size per device; if this is specified, --global_bs will be ignored
+    vae_local_bs: int = 64 # sub-batch size for vae loss calculation
+    global_bs: int = 0  # global batch size (exclusive to --local_bs)
+    lr: float = 5e-4  # learning rate
+    wd: float = 0.02  # weight decay
+    disc_lr: float = 2e-5  # disc lr
+    disc_wd: float = 0.2
+    grad_clip: float = 10  # <=0 for not using grad clip
+    ema: float = 0.9999  # ema ratio
+    warmup_iter: int = None
+    warmup_ep: float = 0.01  # lr warmup: epochs
+    disc_start_ep: float = 0.375  # start using disc loss for VAE after xxx epochs;
+    disc_warmup_ep: float = 0.03  # disc loss warm up epochs;
+    schedule: str = 'cos'  # lr schedule type
+    lr_start_ratio: float = 0.  # lr warmup: initial lr ratio
+    lr_end_ratio: float = 0.1  # lr schedule: final lr ratio
+    disc_lr_end_ratio: float = 0.1
+    custom_lr_multiplier: float = None
+    optimizer: str = 'adamw'
+    optim_eps: float = 1e-6
+    fuse_opt: bool = False  # whether to use fused optimizer
+    optim_beta: str = '0.9_0.95'  # beta1, beta2 of optimizer
+    disc_optim_beta: str = '0.5_0.9'  # beta1, beta2 of disc optimizer
+
+    # loss
+    l1: float = 0.2  # L1 rec loss weight
+    l2: float = 1.0  # L2 rec loss weight
+    lp: float = 1.0  # lpips loss weight
+    lpr: int = 48    # only calculate lpips >= this image resolution
+    ld: float = 0.4  # discriminator loss weight; if <0: NO ADAPTIVE WEIGHT
+    le: float = 0.0  # VQ entropy loss weight
+    lq: float = 1.0
+    lc: float = 1.0  # CLIP loss weight
+    e_temp: float = 0.01
+    gada: int = 1
+    bcr: float = 4.  # balanced Consistency Regularization, used on small dataset with low reso, StyleSwin: 10.0
+    bcr_cut: float = 0.2  # cutout ratio (0.5: 50% width)
+    dcrit: str = 'hg'  # hg hinge, sp softplus, ln linear
+
+    # wandb log
+    report_wandb: bool = True
+    wandb_notes: str = None
+    run_id: str = None
+
+    # debug
+    eval_per_epoch: int = 8
+    dbg_unused_param: bool = False
+    dbg_nan: bool = False  # 'KEVIN_LOCAL' in os.environ
+    seed: int = None
+    deterministic: bool = False
+    same_seed_for_all_ranks: int = 0  # this is only for distributed sampler
+
+
+    def seed_everything(self):
+        torch.backends.cudnn.enabled = True
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
+        if self.seed is not None:
+            if self.deterministic:
+                torch.backends.cudnn.benchmark = False
+                torch.backends.cudnn.deterministic = True
+                torch.use_deterministic_algorithms(True)
+                os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8'
+            seed = self.seed + dist.get_rank() * 10000
+            os.environ['PYTHONHASHSEED'] = str(seed)
+            random.seed(seed)
+            np.random.seed(seed)
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+
+    def get_different_generator_for_each_rank(self) -> Optional[torch.Generator]:  # for random augmentation
+        if self.seed is None:
+            return None
+        g = torch.Generator()
+        g.manual_seed(self.seed * dist.get_world_size() + dist.get_rank())
+        return g
+
+    def state_dict(self, key_ordered=True) -> Union[OrderedDict, dict]:
+        d = (OrderedDict if key_ordered else dict)()
+        for k in self.class_variables.keys():
+            if k not in {'device'}:  # these are not serializable
+                d[k] = getattr(self, k)
+        return d
+
+    def load_state_dict(self, state_dict):
+        for k, v in state_dict.items():
+            try:
+                setattr(self, k, v)
+            except Exception as e:
+                print(f'k={k}, v={v}')
+                raise e
+
+    @staticmethod
+    def set_tf32(tf32: bool):
+        if torch.cuda.is_available():
+            torch.backends.cudnn.allow_tf32 = bool(tf32)
+            torch.backends.cuda.matmul.allow_tf32 = bool(tf32)
+            if hasattr(torch, 'set_float32_matmul_precision'):
+                torch.set_float32_matmul_precision('high' if tf32 else 'highest')
+                print(f'[tf32] [precis] torch.get_float32_matmul_precision(): {torch.get_float32_matmul_precision()}')
+            print(f'[tf32] [ conv ] torch.backends.cudnn.allow_tf32: {torch.backends.cudnn.allow_tf32}')
+            print(f'[tf32] [matmul] torch.backends.cuda.matmul.allow_tf32: {torch.backends.cuda.matmul.allow_tf32}')
+
+    def __str__(self):
+        s = []
+        for k in self.class_variables.keys():
+            if k not in {'device', 'dbg_ks_fp'}:  # these are not serializable
+                s.append(f'  {k:20s}: {getattr(self, k)}')
+        s = '\n'.join(s)
+        return f'{{\n{s}\n}}\n'
+
+
+def init_dist_and_get_args():
+    for i in range(len(sys.argv)):
+        if sys.argv[i].startswith('--local-rank=') or sys.argv[i].startswith('--local_rank='):
+            del sys.argv[i]
+            break
+
+    args = Args(explicit_bool=True).parse_args(known_only=True)
+    # warn args.extra_args
+    if len(args.extra_args) > 0:
+        print(f'======================================================================================')
+        print(f'=========================== WARNING: UNEXPECTED EXTRA ARGS ===========================\n{args.extra_args}')
+        print(f'=========================== WARNING: UNEXPECTED EXTRA ARGS ===========================')
+        print(f'======================================================================================\n\n')
+
+    # init torch distributed
+    os.makedirs(args.output_dir, exist_ok=True)
+    dist.init_distributed_mode(local_out_path=args.output_dir, timeout_minutes=30)
+
+    # set env
+    args.set_tf32(args.tf32)
+    args.seed_everything()
+    args.device = dist.get_device()
+
+    # update args
+    if args.local_bs == 0:
+        args.local_bs = max(1, round(args.global_bs / args.grad_accu / dist.get_world_size()))
+    args.global_bs = args.local_bs * dist.get_world_size()
+    if args.fp16 or args.bf16:
+        args.dtype = torch.float16 if args.fp16 else torch.bfloat16
+
+    return args

native/utils/dist.py

@@ -0,0 +1,302 @@
+import datetime
+import functools
+import os
+import sys
+from typing import List
+from typing import Union
+
+import pytz
+import torch
+import torch.distributed as tdist
+import torch.multiprocessing as mp
+
+__rank, __local_rank, __world_size, __device = 0, 0, 1, 'cuda' if torch.cuda.is_available() else 'cpu'
+__rank_str_zfill = '0'
+__initialized = False
+
+
+def initialized():
+    return __initialized
+
+
+def __initialize(fork=False, backend='nccl', gpu_id_if_not_distibuted=0, timeout_minutes=30):
+    global __device
+    if not torch.cuda.is_available():
+        print(f'[dist initialize] cuda is not available, use cpu instead', file=sys.stderr)
+        return
+    elif 'RANK' not in os.environ:
+        torch.cuda.set_device(gpu_id_if_not_distibuted)
+        __device = torch.empty(1).cuda().device
+        print(f'[dist initialize] env variable "RANK" is not set, use {__device} as the device', file=sys.stderr)
+        return
+    # then 'RANK' must exist
+    global_rank, num_gpus = int(os.environ['RANK']), torch.cuda.device_count()
+    local_rank = global_rank % num_gpus
+    torch.cuda.set_device(local_rank)
+    
+    # ref: https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py#L29
+    if mp.get_start_method(allow_none=True) is None:
+        method = 'fork' if fork else 'spawn'
+        print(f'[dist initialize] mp method={method}')
+        mp.set_start_method(method)
+    tdist.init_process_group(backend=backend, timeout=datetime.timedelta(seconds=timeout_minutes * 60))
+    
+    global __rank, __local_rank, __world_size, __initialized, __rank_str_zfill
+    __local_rank = local_rank
+    __rank, __world_size = tdist.get_rank(), tdist.get_world_size()
+    __rank_str_zfill = str(__rank).zfill(len(str(__world_size)))
+    __device = torch.empty(1).cuda().device
+    __initialized = True
+    
+    assert tdist.is_initialized(), 'torch.distributed is not initialized!'
+    print(f'[lrk={get_local_rank()}, rk={get_rank()}]')
+
+
+def get_rank():
+    return __rank
+
+
+def get_rank_str_zfill():
+    return __rank_str_zfill
+
+
+def get_local_rank():
+    return __local_rank
+
+
+def get_world_size():
+    return __world_size
+
+
+def get_device():
+    return __device
+
+
+def set_gpu_id(gpu_id: int):
+    if gpu_id is None: return
+    global __device
+    if isinstance(gpu_id, (str, int)):
+        torch.cuda.set_device(int(gpu_id))
+        __device = torch.empty(1).cuda().device
+    else:
+        raise NotImplementedError
+
+
+def is_master():
+    return __rank == 0
+
+
+def is_local_master():
+    return __local_rank == 0
+
+
+def new_group(ranks: List[int]):
+    if __initialized:
+        return tdist.new_group(ranks=ranks)
+    return None
+
+
+def new_local_machine_group():
+    if __initialized:
+        cur_subgroup, subgroups = tdist.new_subgroups()
+        return cur_subgroup
+    return None
+
+
+def barrier():
+    if __initialized:
+        tdist.barrier()
+
+
+def allreduce(t: torch.Tensor, async_op=False):
+    if __initialized:
+        if not t.is_cuda:
+            cu = t.detach().cuda()
+            ret = tdist.all_reduce(cu, async_op=async_op)
+            t.copy_(cu.cpu())
+        else:
+            ret = tdist.all_reduce(t, async_op=async_op)
+        return ret
+    return None
+
+
+def allgather(t: torch.Tensor, cat=True) -> Union[List[torch.Tensor], torch.Tensor]:
+    if __initialized:
+        if not t.is_cuda:
+            t = t.cuda()
+        ls = [torch.empty_like(t) for _ in range(__world_size)]
+        tdist.all_gather(ls, t)
+    else:
+        ls = [t]
+    if cat:
+        ls = torch.cat(ls, dim=0)
+    return ls
+
+
+def allgather_diff_shape(t: torch.Tensor, cat=True) -> Union[List[torch.Tensor], torch.Tensor]:
+    if __initialized:
+        if not t.is_cuda:
+            t = t.cuda()
+        
+        t_size = torch.tensor(t.size(), device=t.device)
+        ls_size = [torch.empty_like(t_size) for _ in range(__world_size)]
+        tdist.all_gather(ls_size, t_size)
+        
+        max_B = max(size[0].item() for size in ls_size)
+        pad = max_B - t_size[0].item()
+        if pad:
+            pad_size = (pad, *t.size()[1:])
+            t = torch.cat((t, t.new_empty(pad_size)), dim=0)
+        
+        ls_padded = [torch.empty_like(t) for _ in range(__world_size)]
+        tdist.all_gather(ls_padded, t)
+        ls = []
+        for t, size in zip(ls_padded, ls_size):
+            ls.append(t[:size[0].item()])
+    else:
+        ls = [t]
+    if cat:
+        ls = torch.cat(ls, dim=0)
+    return ls
+
+
+def broadcast(t: torch.Tensor, src_rank) -> None:
+    if __initialized:
+        if not t.is_cuda:
+            cu = t.detach().cuda()
+            tdist.broadcast(cu, src=src_rank)
+            t.copy_(cu.cpu())
+        else:
+            tdist.broadcast(t, src=src_rank)
+
+
+def dist_fmt_vals(val: float, fmt: Union[str, None] = '%.2f') -> Union[torch.Tensor, List]:
+    if not initialized():
+        return torch.tensor([val]) if fmt is None else [fmt % val]
+    
+    ts = torch.zeros(__world_size)
+    ts[__rank] = val
+    allreduce(ts)
+    if fmt is None:
+        return ts
+    return [fmt % v for v in ts.cpu().numpy().tolist()]
+
+
+def master_only(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if force or is_master():
+            ret = func(*args, **kwargs)
+        else:
+            ret = None
+        barrier()
+        return ret
+    return wrapper
+
+
+def local_master_only(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if force or is_local_master():
+            ret = func(*args, **kwargs)
+        else:
+            ret = None
+        barrier()
+        return ret
+    return wrapper
+
+
+def for_visualize(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if is_master():
+            # with torch.no_grad():
+            ret = func(*args, **kwargs)
+        else:
+            ret = None
+        return ret
+    return wrapper
+
+
+def finalize():
+    if __initialized:
+        tdist.destroy_process_group()
+
+
+def init_distributed_mode(local_out_path, only_sync_master=False, timeout_minutes=30):
+    try:
+        __initialize(fork=False, timeout_minutes=timeout_minutes)
+        barrier()
+    except RuntimeError as e:
+        print(f'{"!"*80}   dist init error (NCCL Error?), stopping training!   {"!"*80}', flush=True)
+        raise e
+    
+    if local_out_path is not None: os.makedirs(local_out_path, exist_ok=True)
+    _change_builtin_print(is_local_master())
+    if (is_master() if only_sync_master else is_local_master()) and local_out_path is not None and len(local_out_path):
+        sys.stdout, sys.stderr = BackupStreamToFile(local_out_path, for_stdout=True), BackupStreamToFile(local_out_path, for_stdout=False)
+
+
+def _change_builtin_print(is_master):
+    import builtins as __builtin__
+    
+    builtin_print = __builtin__.print
+    if type(builtin_print) != type(open):
+        return
+    
+    def prt(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        clean = kwargs.pop('clean', False)
+        deeper = kwargs.pop('deeper', False)
+        if is_master or force:
+            if not clean:
+                f_back = sys._getframe().f_back
+                if deeper and f_back.f_back is not None:
+                    f_back = f_back.f_back
+                file_desc = f'{f_back.f_code.co_filename:24s}'[-24:]
+                time_str = datetime.datetime.now(tz=pytz.timezone('Asia/Shanghai')).strftime('[%m-%d %H:%M:%S]')
+                builtin_print(f'{time_str} ({file_desc}, line{f_back.f_lineno:-4d})=>', *args, **kwargs)
+            else:
+                builtin_print(*args, **kwargs)
+    
+    __builtin__.print = prt
+
+
+class BackupStreamToFile(object):
+    def __init__(self, local_output_dir, for_stdout=True):
+        self.for_stdout = for_stdout
+        self.terminal_stream = sys.stdout if for_stdout else sys.stderr
+        fname = os.path.join(local_output_dir, 'backup1_stdout.txt' if for_stdout else 'backup2_stderr.txt')
+        existing = os.path.exists(fname)
+        self.file_stream = open(fname, 'a')
+        if existing:
+            time_str = datetime.datetime.now(tz=pytz.timezone('Asia/Shanghai')).strftime('[%m-%d %H:%M:%S]')
+            self.file_stream.write('\n'*7 + '='*55 + f'   RESTART {time_str}   ' + '='*55 + '\n')
+        self.file_stream.flush()
+        self.enabled = True
+    
+    def write(self, message):
+        self.terminal_stream.write(message)
+        self.file_stream.write(message)
+    
+    def flush(self):
+        self.terminal_stream.flush()
+        self.file_stream.flush()
+    
+    def close(self):
+        if not self.enabled:
+            return
+        self.enabled = False
+        self.file_stream.flush()
+        self.file_stream.close()
+        if self.for_stdout:
+            sys.stdout = self.terminal_stream
+            sys.stdout.flush()
+        else:
+            sys.stderr = self.terminal_stream
+            sys.stderr.flush()
+    
+    def __del__(self):
+        self.close()

preprocessor_config.json

@@ -0,0 +1,12 @@
+{
+  "reference": "https://github.com/FoundationVision/UniTok/blob/bb8012e8e8d7be66928bedd80506c56a55ea6305/utils/data.py",
+  "reference_function": "build_transforms: preprocess_val",
+  "resize_shortest_edge": 256,
+  "resize_interpolation": "PIL Lanczos",
+  "center_crop": [256, 256],
+  "convert_to_rgb": true,
+  "to_tensor_scale": "uint8 / 255",
+  "normalization": "tensor * 2 - 1",
+  "layout": "NCHW",
+  "notes": "This describes the native deterministic validation transform, not every downstream MLLM prompt processor."
+}

release_manifest.json

@@ -0,0 +1,117 @@
+{
+  "as_of": "2026-10-21",
+  "files": {
+    "model.safetensors": {
+      "size": 3281477448,
+      "sha256": "364bbd3e3f453f19dba447805f347aabdd7abdc7da20186364f5b2961c541b5c"
+    },
+    "tensor_mapping.json": {
+      "size": 247482,
+      "sha256": "b1bd76fbe4fd06157d4e6384d619c604b7bb498decf756fb2d76f79226dbffe8"
+    },
+    "OPENCLIP_LICENSE": {
+      "size": 1195,
+      "sha256": "d4a357b7ebcb1eec1502d4ee0dc79c73469e126937cf66d5df4c20a9319b5217"
+    },
+    "code_provenance.json": {
+      "size": 3602,
+      "sha256": "5592e18e7f3e1dcd5c5844047c5ee80235d1ae78e642cb3cf10f9052601a3c25"
+    },
+    "LICENSE": {
+      "size": 11540,
+      "sha256": "76510cc65779339b0a4d3a510200eb18e8dd5c79dc3ad7958530da499160ad2c"
+    },
+    "requirements.txt": {
+      "size": 93,
+      "sha256": "6820fdfb476caef1b0553cb1929da16e2ee19caad3c32fa623b632edbe3d29d2"
+    },
+    "config.json": {
+      "size": 2590,
+      "sha256": "4713288a6bcc010b0977fa6bd5f95ee5c0a23d283cdc248afcba08dccf8344c2"
+    },
+    "tensor_sha256.json": {
+      "size": 143621,
+      "sha256": "4386911017688407b41fb0e295846a3db213a799405256eb56e5b519bf22cf26"
+    },
+    "README.md": {
+      "size": 3631,
+      "sha256": "023fbaf021040ff8f7d95ec4a5dabb048df2fbe4e48e5f004893f1faae68da24"
+    },
+    "validate.py": {
+      "size": 3551,
+      "sha256": "7ed5f8c436d9fbb602b84a763c4c3d131ca03ca97c5dfa5dc1bceb4156572379"
+    },
+    "validation.json": {
+      "size": 1995,
+      "sha256": "c7465f3e73c9ce27928fe9c2dc0e7460ce3c9b30e2d4b08715c599a6e000ff05"
+    },
+    "CODE_LICENSE": {
+      "size": 1072,
+      "sha256": "a318f4da9d043e05f535b65774febde7572908c97562cc14e49a2f3687f5a453"
+    },
+    "extraction.json": {
+      "size": 953,
+      "sha256": "f689cf9bb9e6f90e65c66fa9306ea15792609d294b2489fe9f5dbeb36204cd2f"
+    },
+    "component.py": {
+      "size": 735,
+      "sha256": "abd48d6e68f4f2e1371c793882e540e24ed68c8ca5ce94632976cf7f032e0d88"
+    },
+    "preprocessor_config.json": {
+      "size": 524,
+      "sha256": "3a8598551509dff8de647325f146892b2e47ed8017e1e707074df9a3dea9a97d"
+    },
+    "native/__init__.py": {
+      "size": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    "native/utils/config.py": {
+      "size": 9535,
+      "sha256": "2776f3ea4baf57708dbaf385568ad26fc10be61b7f1b01e8a9d897b3d4242aed"
+    },
+    "native/utils/__init__.py": {
+      "size": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    "native/utils/dist.py": {
+      "size": 8957,
+      "sha256": "4ad98f72e600a666795eb074329f4e6cfd8c182776fb7f9e4d9d39ae2dd83524"
+    },
+    "native/models/unitok.py": {
+      "size": 7009,
+      "sha256": "bc18a150c277e85a052bf03ed09cc6c4ecc51745cd926cb680b4d352d03fc0ac"
+    },
+    "native/models/__init__.py": {
+      "size": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    "native/models/vitamin.py": {
+      "size": 27654,
+      "sha256": "bd10601ae44739f7c67273c1f2c783e9205ecc23440f63b6dc50fdd4abae896f"
+    },
+    "native/models/vqvae.py": {
+      "size": 6425,
+      "sha256": "1afad9ee7f0624c1af6265c245c54d89d9a732f58fa3b2c5c2c8db7d9e00fdc3"
+    },
+    "native/models/quant.py": {
+      "size": 7558,
+      "sha256": "fcd9fa6a76a962eab32c33a7deb4205dfaac019f8816be41951f0604e2343b3b"
+    },
+    "native/open_clip/source_model.py": {
+      "size": 23383,
+      "sha256": "2a7bc461f7cfe70497e7ce8541f75a8b67dcbd3d93685cfd8671ebf98d337150"
+    },
+    "native/open_clip/__init__.py": {
+      "size": 0,
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+    },
+    "native/open_clip/utils.py": {
+      "size": 3472,
+      "sha256": "f0267bb9d5402b1b5b38a90f96f6cd9002ac21d2db0f3f3fbee8ed48b2f46047"
+    },
+    "native/open_clip/transformer.py": {
+      "size": 27556,
+      "sha256": "dfaef403435d5766aeea8286fca0f9fc46285ee25fd28e87ba197bf51f1f676e"
+    }
+  }
+}

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.2
+torchvision
+safetensors>=0.4
+timm==1.0.8
+typed-argument-parser==1.12.0
+pytz
+numpy

validate.py

@@ -0,0 +1,48 @@
+import hashlib, json, time
+from pathlib import Path
+import torch
+from importlib.metadata import version
+torch.set_num_threads(4)
+from safetensors import safe_open
+from component import load_model
+
+root=Path(__file__).parent
+started=time.monotonic(); checks=[]
+def check(name,ok,**details):
+    assert ok,(name,details)
+    checks.append({'name':name,'passed':True,**details});print(name,details,flush=True)
+
+hashes=json.loads((root/'tensor_sha256.json').read_text())
+with safe_open(root/'model.safetensors',framework='pt') as f:
+    assert set(f.keys())==set(hashes)
+    for k in f.keys():
+        t=f.get_tensor(k)
+        assert hashlib.sha256(t.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()).hexdigest()==hashes[k],k
+check('1335_source_tensor_payloads',len(hashes)==1335)
+m=load_model(root)
+check('native_strict_load',len(m.state_dict())==1335)
+assert type(m.text_encoder).__module__=='native.open_clip.transformer'
+check('eight_4096_by_8_codebooks',len(m.quantizer.codebooks)==8 and all(tuple(c.codebook.weight.shape)==(4096,8) for c in m.quantizer.codebooks))
+with torch.inference_mode():
+    x=torch.linspace(-1,1,3*256*256).reshape(1,3,256,256)
+    ids=m.img_to_idx(x)
+    check('token_geometry',tuple(ids.shape)==(1,8,256),shape=list(ids.shape))
+    check('per_codebook_ID_bounds',ids.min().item()>=0 and ids.max().item()<4096)
+    q=m.quantizer.idx_to_f(ids)
+    canonical=torch.cat([c.codebook(ids[:,i]) for i,c in enumerate(m.quantizer.codebooks)],dim=-1)
+    check('normalized_lookup_exact',torch.equal(q,canonical))
+    rgb=m.idx_to_img(ids)
+    direct=m.decoder(m.post_quant_proj(canonical)).clamp(-1,1)
+    check('decode_ID_lookup_exact',torch.equal(rgb,direct))
+    check('finite_RGB_reconstruction',rgb.shape==x.shape and torch.isfinite(rgb).all().item(),mse=(rgb-x).square().mean().item())
+    reconstructed=m.img_to_reconstructed_img(x)
+    check('native_quantizer_reconstruction_bound',torch.allclose(reconstructed,rgb,atol=2e-5,rtol=2e-5),max_abs=(rgb-reconstructed).abs().max().item())
+    visual=m.encode_image(x,normalize=True)
+    text=torch.zeros((1,m.context_length),dtype=torch.long)
+    text[0,0]=49406;text[0,1]=49407
+    textual=m.encode_text(text,normalize=True)
+    check('semantic_image_features',tuple(visual.shape)==(1,768) and torch.isfinite(visual).all().item() and torch.allclose(visual.norm(dim=-1),torch.ones(1),atol=1e-5))
+    check('semantic_text_features',tuple(textual.shape)==(1,768) and torch.isfinite(textual).all().item() and torch.allclose(textual.norm(dim=-1),torch.ones(1),atol=1e-5))
+report={'strict_load':True,'source_tensor_bytes_equal':True,'checks':checks,'smoke':{'passed':True,'device':'cpu','dtype':'float32','resolution':[256,256]},'environment':{p:version(p) for p in ['torch','torchvision','timm','safetensors','typed-argument-parser','pytz']},'text_implementation':'Unchanged official project-vendored native.open_clip.transformer.TextTransformer, same configured default constructor branch. External OpenCLIP is not imported.','elapsed_seconds':time.monotonic()-started,'limitations':['Synthetic native codec and semantic feature fixtures; no downstream image quality, retrieval, VQA or text-to-image benchmark.','This is the complete trained tokenizer mapping including its semantic text tower; no language generator is included.','The eight codebooks each have 4096 entries of width 8. The sum of entries is 32768; IDs are local to each codebook.']}
+(root/'validation.json').write_text(json.dumps(report,indent=2)+'\n')
+print('UNITOK VALIDATED',len(checks),flush=True)

validation.json

@@ -0,0 +1,81 @@
+{
+  "strict_load": true,
+  "source_tensor_bytes_equal": true,
+  "checks": [
+    {
+      "name": "1335_source_tensor_payloads",
+      "passed": true
+    },
+    {
+      "name": "native_strict_load",
+      "passed": true
+    },
+    {
+      "name": "eight_4096_by_8_codebooks",
+      "passed": true
+    },
+    {
+      "name": "token_geometry",
+      "passed": true,
+      "shape": [
+        1,
+        8,
+        256
+      ]
+    },
+    {
+      "name": "per_codebook_ID_bounds",
+      "passed": true
+    },
+    {
+      "name": "normalized_lookup_exact",
+      "passed": true
+    },
+    {
+      "name": "decode_ID_lookup_exact",
+      "passed": true
+    },
+    {
+      "name": "finite_RGB_reconstruction",
+      "passed": true,
+      "mse": 0.0007349226507358253
+    },
+    {
+      "name": "native_quantizer_reconstruction_bound",
+      "passed": true,
+      "max_abs": 1.0728836059570312e-06
+    },
+    {
+      "name": "semantic_image_features",
+      "passed": true
+    },
+    {
+      "name": "semantic_text_features",
+      "passed": true
+    }
+  ],
+  "smoke": {
+    "passed": true,
+    "device": "cpu",
+    "dtype": "float32",
+    "resolution": [
+      256,
+      256
+    ]
+  },
+  "environment": {
+    "torch": "2.14.1",
+    "torchvision": "0.29.1",
+    "timm": "1.0.8",
+    "safetensors": "0.8.0",
+    "typed-argument-parser": "1.12.0",
+    "pytz": "2026.5"
+  },
+  "text_implementation": "Unchanged official project-vendored native.open_clip.transformer.TextTransformer, same configured default constructor branch. External OpenCLIP is not imported.",
+  "elapsed_seconds": 90.44274949999817,
+  "limitations": [
+    "Synthetic native codec and semantic feature fixtures; no downstream image quality, retrieval, VQA or text-to-image benchmark.",
+    "This is the complete trained tokenizer mapping including its semantic text tower; no language generator is included.",
+    "The eight codebooks each have 4096 entries of width 8. The sum of entries is 32768; IDs are local to each codebook."
+  ]
+}