| { | |
| "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 | |
| }, | |
| "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": 1152 | |
| }, | |
| "dec_registers_posemb_module": { | |
| "_target_": "flextok.model.utils.posembs.PositionalEmbeddingAdder", | |
| "read_key": "dec_registers_proj", | |
| "write_key": "dec_registers_proj", | |
| "dim": 1152, | |
| "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": 1152, | |
| "size_sampling_mode": "pow2" | |
| }, | |
| "dec_latent_dropout": { | |
| "_target_": "flextok.model.preprocessors.nullcond.LearnedNullCond", | |
| "read_write_key": "dec_registers_proj", | |
| "dim": 1152, | |
| "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": 1152, | |
| "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": 1152, | |
| "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": 1152, | |
| "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": 1152, | |
| "depth": 18, | |
| "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": false | |
| }, | |
| "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": 1152, | |
| "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 | |
| } | |
| } | |
| } | |
| }, | |
| "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 | |
| }, | |
| "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 | |
| }, | |
| "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 | |
| }, | |
| "_target_": "flextok.flextok_wrapper.FlexTok", | |
| "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": false | |
| }, | |
| "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 | |
| } | |
| } | |
| } | |
| } |