[content] genie.py (python)
    def restore_from_old_format(checkpoint_path):
        """Restore from old PyTreeCheckpointer format"""
        dummy_tokenizer = TokenizerVQVAE(
            in_dim=args.image_channels,
            model_dim=args.tokenizer_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_patch_latents,
            patch_size=args.patch_size,
            num_blocks=args.tokenizer_num_blocks,
            num_heads=args.tokenizer_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
        )
        tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
        dummy_tokenizer_train_state = TrainState.create(
            apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
        )
        abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
            dummy_tokenizer_train_state, sharding
        )
        tokenizer_restore_target = {"model": abstract_sharded_tokenizer_state}
        tokenizer_restore_args = orbax_utils.restore_args_from_target(
            tokenizer_restore_target
        )
        restored_tokenizer_params = (
            PyTreeCheckpointer()
            .restore(
                checkpoint_path,
                item=tokenizer_restore_target,
                restore_args=tokenizer_restore_args,
            )["model"]
            .params["params"]
        )
        return restored_tokenizer_params

    def restore_from_new_format(checkpoint_path):
        """Restore from new composite checkpoint manager format"""
        import orbax.checkpoint as ocp
        
        # Create a dummy train state to get the structure
        dummy_tokenizer = TokenizerVQVAE(
            in_dim=args.image_channels,
            model_dim=args.tokenizer_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_patch_latents,
            patch_size=args.patch_size,
            num_blocks=args.tokenizer_num_blocks,
            num_heads=args.tokenizer_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
        )
        tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
        dummy_tokenizer_train_state = TrainState.create(
            apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
        )
        
        # Create checkpoint manager to restore
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add('model_state', ocp.args.StandardRestore, ocp.handlers.StandardCheckpointHandler)
        
        checkpoint_manager = ocp.CheckpointManager(
            os.path.dirname(checkpoint_path),
            options=ocp.CheckpointManagerOptions(),
            handler_registry=handler_registry,
        )
        
        # Get the step from the checkpoint path
        checkpoint_name = os.path.basename(checkpoint_path)
        step = int(checkpoint_name.split('_')[-1]) if '_' in checkpoint_name else 0
        
        abstract_train_state = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, dummy_tokenizer_train_state)
        restored = checkpoint_manager.restore(
            step,
            args=ocp.args.StandardRestore(abstract_train_state),
        )
        return restored["model_state"].params["params"]

    # Restore tokenizer
    if args.tokenizer_checkpoint:
        try:
            # Try new format first
            restored_tokenizer_params = restore_from_new_format(args.tokenizer_checkpoint)
        except Exception:
            # Fall back to old format
            restored_tokenizer_params = restore_from_old_format(args.tokenizer_checkpoint)
        
        train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)