unable to load using pipe = pipeline("image-classification", model= MODEL_REPO_ID) my custom class derived from  Dinov2ForImageClassification
Author: yuragenetika
Pull request: no

### System Info

- `transformers` version: 4.37.2
- Platform: macOS-14.2.1-arm64-arm-64bit
- Python version: 3.11.7
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.26.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.2.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. Create new class for classification based on dinov2

import torch
import torch.nn as nn
from transformers import Dinov2ForImageClassification
from transformers.modeling_outputs import ImageClassifierOutput

class CustomDinov2Classifier(Dinov2ForImageClassification):
    def __init__(self, config):
        super().__init__(config)
        self.classifier = self.classifier = nn.Sequential(
            nn.Linear(config.hidden_size, config.hidden_size),
            nn.ReLU(),
            nn.Linear(config.hidden_size, self.num_labels)
        )

    def forward(self, pixel_values, output_hidden_states=False, output_attentions=False, labels=None):
        # use frozen features
        outputs = self.dinov2(pixel_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions)
        logits = outputs.last_hidden_state[:,0,:]
        logits = self.classifier(logits)
    
        loss = None
        if labels is not None:
            criterion = torch.nn.CrossEntropyLoss()
            loss = criterion(logits, labels)

        return ImageClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)


2. instantiate this class

id2label = {0: 'CLASS A', 1: 'CLASS B', 2: 'CLASS C', 3: 'CLASS D'}
model = CustomDinov2Classifier.from_pretrained('facebook/dinov2-base-imagenet1k-1-layer', id2label=id2label, label2id=label2id, num_labels=len(id2label), ignore_mismatched_sizes = True)


3. Train this model 
4.  save it with trainer.save_model() or just model.save_pretrained()
5. push this model to hub:

model.push_to_hub(MODEL_REPO_ID)
trainer.push_to_hub(MODEL_REPO_ID)

see that under how to use it section:

from transformers import AutoImageProcessor, CustomDinov2Classifier

processor = AutoImageProcessor.from_pretrained(MODEL_REPO_ID)
classifier = CustomDinov2Classifier.from_pretrained(MODEL_REPO_ID).to(device)

basically the custom class is presented as part of transformers library

additionally its impossible to use trained model without classifying the new class name 

from transformers import pipeline

pipe = pipeline("image-classification", model= MODEL_REPO_ID)

it loads Dinov2ForImageClassification instead of CustomDinov2Classifier

### Expected behavior

I expect that from transformers import pipeline

pipe = pipeline("image-classification", model= MODEL_REPO_ID)

should load my custom class Dinov2ForImageClassification and not Dinov2ForImageClassification

Comments:
Hi @yuragenetika, thanks for raising this issue! 

In order for you to be able to use your custom class in pipelines, you have to register it with its corresponding auto class, such that the following would load your model:

```py
from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("my_checkpoint")
```

You can find instructions on this [here in the docs ](https://huggingface.co/docs/transformers/custom_models#registering-a-model-with-custom-code-to-the-auto-classes)

@yuragenetika, do you mind sharing your repo_id, I'll take a look

in case that `YuraGenetika/roi-classifier-20x-GFP-scenes_large-mpl-weights_label_smoothing` is your repo, you need to add your custom architecture to the hub first, this problem is not related to pipeline at all.
as for the model kindly confim if that is your repo, and I'll fix it for you.
you might also consider reading this blogpost about using custom architecture when working with huggingface : https://huggingface.co/blog/not-lain/custom-architectures-with-huggingface

Thanks for the response.
I followed exactly the instructions and I get the following error on reading the model from hf remote repository:

from transformers import AutoImageProcessor, AutoModelForImageClassification

processor = AutoImageProcessor.from_pretrained(MODEL_REPO_ID, trust_remote_code=True)
model = AutoModelForImageClassification.from_pretrained(MODEL_REPO_ID, trust_remote_code=True)

OSError: YuraGenetika/sampler-test does not appear to have a file named preprocessor_config.json. Checkout 'https://huggingface.co/YuraGenetika/sampler-test/main' for available files.

Another question,
I am basically want to overrider classifier head of Dinov2ForImageClassification class but the instructions in above don't provide clear description how to load appropriate weights for dinov2 different sizes

for example I want to reuse wights from facebook/dinov2-large-imagenet1k-1-layer

@yuragenetika did you setup your repos to private ? 

![image](https://github.com/huggingface/transformers/assets/70411813/3caad8c3-a14c-4e58-810e-442ace58200b)


since you changed the architecture for the model you need to add the new architecture to your repo by adding the following files: 
* `config.py` 
* `custom_model_architecture.py` 
these 2 files need to be defined in orfer for the transformers library to download your weights and instantiate your model 

again, the transformers library cannot load a custom architecture unless it was ___defined in the hub___ and the `config.json` is pointing at it. defining the architecture in __main__ does not fix this.
try checking these resources to find out more about custom architectures : 
* https://huggingface.co/blog/not-lain/custom-architectures-with-huggingface
* https://huggingface.co/docs/transformers/custom_models
* https://huggingface.co/docs/transformers/en/add_new_pipeline

there is another work around this, but not recommended since there is too much manual changes in this : 
```python
(...)

mymodel = CustomDinov2Classifier(...)
model.load( 
   # load weights
)

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("image-classification", model="facebook/dinov2-small-imagenet1k-1-layer")

# override pipeline parameters
pipe.model = mymodel
pipe.model.config.label2id = ....
pipe.model.config.id2label = ...

pipe(new_input)
(...)
```

but this is **NOT RECOMMENDED METHOD**


thanks for your help

what really worked for me and was the following:

BACKBONE_REPO_ID = 'facebook/dinov2-small-imagenet1k-1-layer'

processor = AutoImageProcessor.from_pretrained(BACKBONE_REPO_ID)
model = CustomDinov2ForImageClassification.from_pretrained(BACKBONE_REPO_ID, id2label=id2label, label2id=label2id, num_labels=len(id2label), ignore_mismatched_sizes = True)
**model.register_for_auto_class("AutoModelForImageClassification")**

no need to redefine config and just define any parameter on class 
        
model.loss_weights=[10.0, 4.0, 2.0, 5.0]
model.label_smoothing = 0.0

train and save

**processor.push_to_hub(MODEL_REPO_ID)**
model.push_to_hub(MODEL_REPO_ID)
trainer.push_to_hub(MODEL_REPO_ID)

after that everything worked as designed


from transformers import pipeline

pipe = pipeline("image-classification", model= MODEL_REPO_ID, trust_remote_code=True)

or 

from transformers import AutoImageProcessor, AutoModelForImageClassification

processor = AutoImageProcessor.from_pretrained(MODEL_REPO_ID, trust_remote_code=True)
classifier = AutoModelForImageClassification.from_pretrained(MODEL_REPO_ID, trust_remote_code=True)





this is my custom model

class CustomDinov2ForImageClassification(Dinov2ForImageClassification):
    # config_class = CustomDinov2Config
    def __init__(self, config):
        super().__init__(config)
        layers = []
        config.num_layers = 1
        input_size, hidden_size = 2 * config.hidden_size, config.hidden_size
        for _ in range(config.num_layers):
            layers.append(nn.Linear(input_size, hidden_size))
            layers.append(nn.GELU())
            layers.append(nn.LayerNorm(config.hidden_size))
            input_size = hidden_size
        self.embeddings = nn.Sequential(*layers)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, pixel_values, output_hidden_states=False, output_attentions=False, labels=None):
        outputs = self.dinov2(
            pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states
        )
        sequence_output = outputs[0]  # batch_size, sequence_length, hidden_size

        cls_token = sequence_output[:, 0]
        patch_tokens = sequence_output[:, 1:]

        linear_input = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1)
        
        embeddings = self.embeddings(linear_input)

        logits = self.classifier(embeddings)
    
        loss = None
        if labels is not None:
            criterion = nn.CrossEntropyLoss(
                weight=torch.tensor(self.loss_weights).to(self.device), 
                label_smoothing=self.label_smoothing)
            loss = criterion(logits, labels)

        return ImageClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)