Initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,79 @@
+---
+license: apple-amlr
+---
+
+# FlexTok: Resampling Images into 1D Token Sequences of Flexible Length
+
+[`Website`](https://flextok.epfl.ch) | [`arXiv`](https://arxiv.org/abs/2502.13967) | [`GitHub`](https://github.com/apple/ml-flextok) | [`🤗 Demo`](https://huggingface.co/spaces/EPFL-VILAB/FlexTok) | [`BibTeX`](#citation)
+
+Official implementation and pre-trained models for: <br>
+[**FlexTok: Resampling Images into 1D Token Sequences of Flexible Length**](https://arxiv.org/abs/2502.13967), arXiv 2025 <br>
+*[Roman Bachmann](https://roman-bachmann.github.io/)\*, [Jesse Allardice](https://github.com/JesseAllardice)\*, [David Mizrahi](https://dmizrahi.com/)\*, [Enrico Fini](https://scholar.google.com/citations?user=OQMtSKIAAAAJ), [Oğuzhan Fatih Kar](https://ofkar.github.io/), [Elmira Amirloo](https://elamirloo.github.io/), [Alaaeldin El-Nouby](https://aelnouby.github.io/), [Amir Zamir](https://vilab.epfl.ch/zamir/), [Afshin Dehghan](https://scholar.google.com/citations?user=wcX-UW4AAAAJ)*
+
+
+## Installation
+For install instructions, please see https://github.com/apple/ml-flextok. 
+
+
+## Usage
+
+To load the `FlexTok d18-d28 ImageNet-1k` model directly from HuggingFace Hub, call:
+```python
+from flextok.flextok_wrapper import FlexTokFromHub
+model = FlexTokFromHub.from_pretrained('EPFL-VILAB/flextok_d18_d28_in1k').eval()
+```
+
+The model can also be loaded by downloading the `model.safetensors` checkpoint in this repository manually and loading it using our helper functions:
+```python
+from hydra.utils import instantiate
+from flextok.utils.checkpoint import load_safetensors
+
+ckpt, config = load_safetensors('/path/to/model.safetensors')
+model = instantiate(config).eval()
+model.load_state_dict(ckpt)
+```
+
+After loading a FlexTok model, image batches can be encoded using:
+```python
+from flextok.utils.demo import imgs_from_urls
+# Load example images of shape (B, 3, 256, 256), normalized to [-1,1]
+imgs = imgs_from_urls(urls=['https://storage.googleapis.com/flextok_site/nb_demo_images/0.png'])
+
+# tokens_list is a list of [1, 256] discrete token sequences
+tokens_list = model.tokenize(imgs)
+```
+
+The list of token sequences can be truncated in a nested fashion:
+```python
+k_keep = 64 # For example, only keep the first 64 out of 256 tokens
+tokens_list = [t[:,:k_keep] for t in tokens_list]
+```
+
+To decode the tokens with FlexTok's rectified flow decoder, call:
+```python
+# tokens_list is a list of [1, l] discrete token sequences, with l <= 256
+# reconst is a [B, 3, 256, 256] tensor, normalized to [-1,1]
+reconst = model.detokenize(
+    tokens_list,
+    timesteps=20, # Number of denoising steps
+    guidance_scale=7.5, # Classifier-free guidance scale
+    perform_norm_guidance=True, # See https://arxiv.org/abs/2410.02416
+)
+```
+
+
+## Citation
+
+If you find this repository helpful, please consider citing our work:
+```
+@article{flextok,
+    title={{FlexTok}: Resampling Images into 1D Token Sequences of Flexible Length},
+    author={Roman Bachmann and Jesse Allardice and David Mizrahi and Enrico Fini and O{\u{g}}uzhan Fatih Kar and Elmira Amirloo and Alaaeldin El-Nouby and Amir Zamir and Afshin Dehghan},
+    journal={arXiv 2025},
+    year={2025},
+}
+```
+
+## License
+
+The model weights in this repository are released under the Apple Model License for Research.

config.json

@@ -0,0 +1,297 @@
+{
+    "vae": {
+        "_target_": "flextok.vae_wrapper.StableDiffusionVAE",
+        "images_read_key": "rgb",
+        "vae_latents_read_key": "vae_latents_reconst",
+        "vae_latents_write_key": "vae_latents",
+        "images_reconst_write_key": "rgb_reconst",
+        "vae_kl_loss_write_key": "kl_loss",
+        "dtype_override": null,
+        "sample_posterior": true,
+        "compile_encode_fn": false,
+        "force_vae_encode": true,
+        "latent_channels": 16,
+        "scaling_factor": 0.88
+    },
+    "regularizer": {
+        "_target_": "flextok.regularizers.quantize_fsq.FSQ",
+        "latents_read_key": "enc_registers",
+        "quants_write_key": "enc_registers_quant",
+        "tokens_write_key": "tokens",
+        "levels": [
+            8,
+            8,
+            8,
+            5,
+            5,
+            5
+        ],
+        "drop_quant_p": 0.0,
+        "packed_call": false
+    },
+    "encoder": {
+        "_target_": "flextok.model.utils.wrappers.SequentialModuleDictWrapper",
+        "module_dict": {
+            "enc_channels_to_last": {
+                "_target_": "flextok.model.utils.dict_ops.PerSampleOp",
+                "read_key": "vae_latents",
+                "write_key": "vae_latents_bhwc",
+                "per_sample_op": {
+                    "_target_": "flextok.model.utils.dict_ops.channels_first_to_last",
+                    "_partial_": true
+                }
+            },
+            "enc_patch_emb": {
+                "_target_": "flextok.model.preprocessors.patching.PatchEmbedder",
+                "input_tensor_list_read_key": "vae_latents_bhwc",
+                "patches_list_write_key": "enc_vae_latents_patched",
+                "n_patches_write_key": "enc_n_patches",
+                "channels_in": 16,
+                "dim": 1152,
+                "patch_sizes": [
+                    2,
+                    2
+                ],
+                "flatten_patches": false
+            },
+            "enc_posemb_module": {
+                "_target_": "flextok.model.utils.posembs.PositionalEmbeddingAdder",
+                "read_key": "enc_vae_latents_patched",
+                "write_key": "enc_vae_latents_patched",
+                "dim": 1152,
+                "max_sizes": [
+                    16,
+                    16
+                ],
+                "posemb_type": "sincos",
+                "posemb_scaling": "absolute"
+            },
+            "enc_register_module": {
+                "_target_": "flextok.model.preprocessors.registers.Registers1D",
+                "input_tensor_list_read_key": "enc_vae_latents_patched",
+                "register_sizes_read_write_key": "register_sizes",
+                "registers_write_key": "enc_registers",
+                "dim": 1152,
+                "n_min": 256,
+                "n_max": 256,
+                "size_sampling_mode": "uniform",
+                "ordering_mode": "nested"
+            },
+            "enc_seq_packer": {
+                "_target_": "flextok.model.preprocessors.flex_seq_packing.BlockWiseSequencePacker",
+                "input_list_read_keys": [
+                    "enc_vae_latents_patched",
+                    "enc_registers"
+                ],
+                "packed_seq_write_key": "enc_packed_seq",
+                "block_mask_write_key": "enc_block_mask",
+                "inner_packed_shapes_write_key": "enc_ps_inner",
+                "outer_packed_shapes_write_key": "enc_ps_outer",
+                "mask_mode": "causal_last",
+                "pad_to_multiple": 128
+            },
+            "enc_transformer": {
+                "_target_": "flextok.model.trunks.transformers.FlexTransformer",
+                "input_seq_read_key": "enc_packed_seq",
+                "output_seq_write_key": "enc_packed_seq",
+                "dim": 1152,
+                "depth": 18,
+                "block_mask_read_key": "enc_block_mask",
+                "use_act_checkpoint": true
+            },
+            "enc_unpacker": {
+                "_target_": "flextok.model.postprocessors.seq_unpacking.SequenceUnpacker",
+                "packed_seq_read_key": "enc_packed_seq",
+                "inner_seq_write_keys": [
+                    "enc_vae_latents_patched",
+                    "enc_registers"
+                ],
+                "inner_packed_shapes_read_key": "enc_ps_inner",
+                "outer_packed_shapes_read_key": "enc_ps_outer"
+            },
+            "enc_to_latents": {
+                "_target_": "flextok.model.postprocessors.heads.LinearHead",
+                "read_key": "enc_registers",
+                "write_key": "enc_registers",
+                "dim": 1152,
+                "dim_out": 6,
+                "use_mup_readout": false,
+                "weight_init_style": "zero",
+                "dtype_override": null
+            }
+        }
+    },
+    "flow_matching_noise_module": {
+        "_target_": "flextok.flow_matching.noise_modules.MinRFNoiseModule",
+        "clean_images_read_key": "vae_latents",
+        "noised_images_write_key": "vae_latents_noised",
+        "noise_write_key": "flow_noise",
+        "timesteps_write_key": "timesteps",
+        "sigmas_write_key": "sigmas",
+        "ln": false,
+        "stratisfied": false,
+        "mode_scale": 0.25
+    },
+    "decoder": {
+        "_target_": "flextok.model.utils.wrappers.SequentialModuleDictWrapper",
+        "module_dict": {
+            "dec_from_latents": {
+                "_target_": "flextok.model.preprocessors.linear.LinearLayer",
+                "read_key": "enc_registers_quant",
+                "write_key": "dec_registers_proj",
+                "dim_in": 6,
+                "dim": 1792
+            },
+            "dec_registers_posemb_module": {
+                "_target_": "flextok.model.utils.posembs.PositionalEmbeddingAdder",
+                "read_key": "dec_registers_proj",
+                "write_key": "dec_registers_proj",
+                "dim": 1792,
+                "max_sizes": [
+                    256
+                ],
+                "posemb_type": "learnable_sum",
+                "posemb_scaling": "absolute"
+            },
+            "dec_nested_dropout": {
+                "_target_": "flextok.model.preprocessors.token_dropout.MaskedNestedDropout",
+                "read_write_key": "dec_registers_proj",
+                "dim": 1792,
+                "size_sampling_mode": "pow2"
+            },
+            "dec_latent_dropout": {
+                "_target_": "flextok.model.preprocessors.nullcond.LearnedNullCond",
+                "read_write_key": "dec_registers_proj",
+                "dim": 1792,
+                "dropout_prob": 0.2
+            },
+            "dec_noise_channels_to_last": {
+                "_target_": "flextok.model.utils.dict_ops.PerSampleOp",
+                "read_key": "vae_latents_noised",
+                "write_key": "vae_latents_noised_bhwc",
+                "per_sample_op": {
+                    "_target_": "flextok.model.utils.dict_ops.channels_first_to_last",
+                    "_partial_": true
+                }
+            },
+            "dec_noise_patch_emb": {
+                "_target_": "flextok.model.preprocessors.patching.PatchEmbedder",
+                "input_tensor_list_read_key": "vae_latents_noised_bhwc",
+                "patches_list_write_key": "vae_latents_noised_patched",
+                "n_patches_write_key": "dec_n_patches",
+                "channels_in": 16,
+                "dim": 1792,
+                "patch_sizes": [
+                    2,
+                    2
+                ],
+                "flatten_patches": false
+            },
+            "dec_patches_posemb_module": {
+                "_target_": "flextok.model.utils.posembs.PositionalEmbeddingAdder",
+                "read_key": "vae_latents_noised_patched",
+                "write_key": "dec_patches",
+                "dim": 1792,
+                "max_sizes": [
+                    16,
+                    16
+                ],
+                "posemb_type": "sincos",
+                "posemb_scaling": "absolute"
+            },
+            "dec_seq_packer": {
+                "_target_": "flextok.model.preprocessors.flex_seq_packing.BlockWiseSequencePacker",
+                "input_list_read_keys": [
+                    "dec_patches",
+                    "dec_registers_proj"
+                ],
+                "packed_seq_write_key": "dec_packed_seq",
+                "block_mask_write_key": "dec_block_mask",
+                "inner_packed_shapes_write_key": "dec_ps_inner",
+                "outer_packed_shapes_write_key": "dec_ps_outer",
+                "emb_packing_fn_write_key": "emb_packing_fn",
+                "mask_mode": "full",
+                "pad_to_multiple": 128,
+                "per_subseq_embs": true
+            },
+            "dec_time_embedder": {
+                "_target_": "flextok.model.preprocessors.time_embedding.TimestepEmbedder",
+                "timesteps_read_key": "timesteps",
+                "time_embedding_write_key": "dec_temb",
+                "dim": 1792,
+                "frequency_embedding_size": 256,
+                "max_timestep": 1000.0
+            },
+            "dec_transformer": {
+                "_target_": "flextok.model.trunks.transformers.FlexTransformer",
+                "input_seq_read_key": "dec_packed_seq",
+                "output_seq_write_key": "dec_packed_seq",
+                "dim": 1792,
+                "depth": 28,
+                "block_mask_read_key": "dec_block_mask",
+                "adaLN_emb_read_key": "dec_temb",
+                "adaLN_packing_fn_read_key": "emb_packing_fn",
+                "adaLN_expansion": 2,
+                "intermediate_layer_write_key": "dec_packed_seq_repa_layer",
+                "intermediate_layers": [
+                    1
+                ],
+                "use_act_checkpoint": true
+            },
+            "dec_unpacker": {
+                "_target_": "flextok.model.postprocessors.seq_unpacking.SequenceUnpacker",
+                "packed_seq_read_key": "dec_packed_seq",
+                "inner_seq_write_keys": [
+                    "dec_patches",
+                    "dec_registers_proj"
+                ],
+                "inner_packed_shapes_read_key": "dec_ps_inner",
+                "outer_packed_shapes_read_key": "dec_ps_outer"
+            },
+            "dec_repa_unpacker": {
+                "_target_": "flextok.model.postprocessors.seq_unpacking.SequenceUnpacker",
+                "packed_seq_read_key": "dec_packed_seq_repa_layer",
+                "inner_seq_write_keys": [
+                    "dec_patches_repa_layer",
+                    "dec_registers_repa_layer"
+                ],
+                "inner_packed_shapes_read_key": "dec_ps_inner",
+                "outer_packed_shapes_read_key": "dec_ps_outer"
+            },
+            "dec_to_patches": {
+                "_target_": "flextok.model.postprocessors.heads.ToPatchesLinearHead",
+                "read_key": "dec_patches",
+                "write_key": "dec_patches",
+                "dim": 1792,
+                "channels_out": 16,
+                "patch_sizes": [
+                    2,
+                    2
+                ],
+                "use_mup_readout": false,
+                "weight_init_style": "zero",
+                "adaLN_emb_read_key": "dec_temb"
+            },
+            "dec_channels_to_first": {
+                "_target_": "flextok.model.utils.dict_ops.PerSampleOp",
+                "read_key": "dec_patches",
+                "write_key": "vae_latents_reconst",
+                "per_sample_op": {
+                    "_target_": "flextok.model.utils.dict_ops.channels_last_to_first",
+                    "_partial_": true
+                }
+            }
+        }
+    },
+    "_target_": "flextok.flextok_wrapper.FlexTok",
+    "pipeline": {
+        "_target_": "flextok.flow_matching.pipelines.MinRFPipeline",
+        "_partial_": true,
+        "target_sizes_read_key": null,
+        "latents_read_key": "enc_registers_quant",
+        "timesteps_read_key": "timesteps",
+        "noised_images_read_key": "vae_latents_noised",
+        "reconst_write_key": "vae_latents_reconst",
+        "out_channels": 16
+    }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25ee13b1c7bffbbff832051bc8269b5a8293f0cebd6d6f9bd0a4e42f49531143
+size 10163625244