ValueError: Unknown layer: Custom>TFViTMainLayer when using a Google transformer model in Streamlit
Author: bluetail14
Pull request: no
Labels: name: bug

### System Info

transformers == 4.21.1
tensorflow == 2.9.1
streamlit ==1.11.1
Windows 10

### Who can help?

@Rocketknight1  @NielsRogge

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```
import streamlit as st
import numpy as np
from PIL import Image 
import tensorflow as tf

st.title("Binary Human Detection Web App")
st.markdown("Is there a human in office space? 🧍")

## Initialize tensorflow model (This can be loaded before anything else)
path_to_model = "C:/Users/myname/Jupiter_Notebooks/Dataset_Thermal_Project/Camera_videos/Saved_models/model_vit.h5"
model_loader = tf.keras.models.load_model(path_to_model)
model_vit = tf.keras.models.Model(model_loader.inputs, model_loader.outputs)

## Preprocess images
def preprocessImage(photo):
    resize_photo = photo.resize((224,224))
    normalized_photo = np.array(resize_photo)/255 # a normalised 2D array                
    reshaped_photo = normalized_photo.reshape(-1, 224, 224, 3)   # to shape as (1, 224, 224, 3)
    return reshaped_photo

uploaded_file = st.sidebar.file_uploader(" ",type=['jpg', 'jpeg'])    

if uploaded_file is not None:
    ## Use a context manager to make sure to close the file!! 
    with Image.open(uploaded_file) as photo:
        tensorflow_image = preprocessImage(photo)
    
    ## Show preprocessed image
    streamlit_widget_image = st.image(tensorflow_image, 'Uploaded Image', use_column_width=True)

## Do prediction
if st.sidebar.button("Click Here to Predict"):
    
    if uploaded_file is None:
        st.sidebar.write("Please upload an Image to Classify")
    else:      
        ## Pass the preprocessed image to the vit model (not the streamlit widget)
        pred_label = model_vit.predict(tensorflow_image)[0]

        ## Print prediction
        st.sidebar.header("ViT model results:") 
        if pred_label > 0.5: st.sidebar.info('Human is detected')
        else: st.sidebar.info('No human is detected')
```
### when I run this  I get the ValueError
ValueError: Unknown layer: Custom>TFViTMainLayer. Please ensure this object is passed to the `custom_objects` argument. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details.
```
#### my model in Tensorflow 
# Base model pre-trained on ImageNet-21k with the 224x224 image resolution
base_model = TFViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
# Freeze base model
base_model.trainable = False
# Create new model
inputs = keras.Input(shape = (3, 224, 224))
x = data_augmentation_vit(inputs)

vit = base_model.vit(inputs)[0]
vit = keras.layers.GlobalAveragePooling1D()(vit)
vit = tf.keras.layers.Dense(256, activation='relu')(vit)
vit = tf. keras.layers.Dropout(0.15)(vit)
outputs = tf.keras.layers.Dense(1, activation='sigmoid', name='outputs')(vit)

model_vit = tf.keras.Model(inputs, outputs)
print(model_vit.summary())

Model: "model_1"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_2 (InputLayer)        [(None, 3, 224, 224)]     0         
                                                                 
 vit (TFViTMainLayer)        TFBaseModelOutputWithPoo  86389248  
                             ling(last_hidden_state=(            
                             None, 197, 768),                    
                              pooler_output=(None, 76            
                             8),                                 
                              hidden_states=None, att            
                             entions=None)                       
                                                                 
 global_average_pooling1d (G  (None, 768)              0         
 lobalAveragePooling1D)                                          
                                                                 
 dense_2 (Dense)             (None, 256)               196864    
                                                                 
 dropout_37 (Dropout)        (None, 256)               0         
                                                                 
 outputs (Dense)             (None, 1)                 257       
                                                                 
=================================================================
model_vit.save("Saved_models/model_vit.h5")
```
### Expected behavior

I have a model, model_vit.h5, trained and saved in Tensorflow based on google's 
vit-base-patch16-224-in21k model.
I expect it to make a prediction in my app like other models.    Yet I am not sure how to register custom object for the TFViTMainLayer in this model.

Comments:
here is the link to model_vit.h5 https://drive.google.com/file/d/1ASXJ6-QVxV7W-rVUV57pUy5sYK1BokZ4/view?usp=sharing

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

try this 

from transformers import TFViTModel


custom_objects = {'TFViTMainLayer': TFViTModel}

with custom_object_scope(custom_objects):
    # Load your models
    vit_model = tf.keras.models.load_model(vit_model_path)




# this solved my issue ,it turns out this is syntax that we need to follow to register our custom layer used