from transformers import PreTrainedModel
from .configuration_segformer_plusplus import SegformerPlusPlusConfig
from .build_model import SegFormer

class SegformerPlusPlusModel(PreTrainedModel):
    config_class = SegformerPlusPlusConfig

    def __init__(self, config: SegformerPlusPlusConfig):
        super().__init__(config)
        self.model = SegFormer(config)

    def forward(self, x, **kwargs):
        return self.model(x)