#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/pp_lcnet/modular_pp_lcnet.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_pp_lcnet.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 The PaddlePaddle Team and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch
import torch.nn as nn
from torch import Tensor

from ...activations import ACT2FN
from ...backbone_utils import BackboneMixin, filter_output_hidden_states
from ...modeling_layers import GradientCheckpointingLayer
from ...modeling_outputs import BackboneOutput, BaseModelOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...processing_utils import Unpack
from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
from ...utils.generic import merge_with_config_defaults
from ...utils.output_capturing import capture_outputs
from .configuration_pp_lcnet import PPLCNetConfig


class PPLCNetConvLayer(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        stride: int = 1,
        activation: str = "hardswish",
        groups: int = 1,
    ):
        super().__init__()
        self.convolution = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=kernel_size // 2,
            bias=False,
            groups=groups,
        )
        self.normalization = nn.BatchNorm2d(out_channels)
        self.activation = ACT2FN[activation] if activation is not None else nn.Identity()

    def forward(self, input: Tensor) -> Tensor:
        hidden_state = self.convolution(input)
        hidden_state = self.normalization(hidden_state)
        hidden_state = self.activation(hidden_state)
        return hidden_state


class PPLCNetDepthwiseSeparableConvLayer(GradientCheckpointingLayer):
    """
    Depthwise Separable Convolution Layer: Depthwise Conv -> SE Module (optional) -> Pointwise Conv
    Core component of lightweight models (e.g., MobileNet, PP-LCNet) that significantly reduces
    the number of parameters and computational cost.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        stride,
        kernel_size,
        use_squeeze_excitation,
        config,
    ):
        super().__init__()
        self.depthwise_convolution = PPLCNetConvLayer(
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=kernel_size,
            stride=stride,
            groups=in_channels,
            activation=config.hidden_act,
        )
        self.squeeze_excitation_module = (
            PPLCNetSqueezeExcitationModule(in_channels, config.reduction) if use_squeeze_excitation else nn.Identity()
        )
        self.pointwise_convolution = PPLCNetConvLayer(
            in_channels=in_channels,
            kernel_size=1,
            out_channels=out_channels,
            stride=1,
            activation=config.hidden_act,
        )

    def forward(self, hidden_state):
        hidden_state = self.depthwise_convolution(hidden_state)
        hidden_state = self.squeeze_excitation_module(hidden_state)
        hidden_state = self.pointwise_convolution(hidden_state)

        return hidden_state


class PPLCNetSqueezeExcitationModule(nn.Module):
    """
    Squeeze-and-Excitation (SE) Module: Adaptive feature recalibration
    Enhances the model's ability to focus on important channels by learning channel-wise attention weights.
    """

    def __init__(self, channel, reduction=4):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)

        self.convolutions = nn.ModuleList()
        for in_channels, out_channels, activation in [
            [channel, channel // reduction, nn.ReLU()],
            [channel // reduction, channel, nn.Hardsigmoid()],
        ]:
            self.convolutions.append(
                nn.Conv2d(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=1,
                    stride=1,
                    padding=0,
                    bias=True,
                )
            )
            self.convolutions.append(activation)

    def forward(self, hidden_state):
        residual = hidden_state
        hidden_state = self.avg_pool(hidden_state)
        for layer in self.convolutions:
            hidden_state = layer(hidden_state)
        hidden_state = residual * hidden_state

        return hidden_state


def make_divisible(value: int, divisor: int = 8, min_value: int | None = None) -> int:
    """
    Ensure that all layers have a channel count that is divisible by `divisor`.
    """
    if min_value is None:
        min_value = divisor
    new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_value < 0.9 * value:
        new_value += divisor
    return int(new_value)


class PPLCNetBlock(nn.Module):
    def __init__(self, config, stage_index):
        super().__init__()
        self.config = config

        blocks = config.block_configs[stage_index]

        self.layers = nn.ModuleList()
        for kernel_size, in_channels, out_channels, stride, use_squeeze_excitation in blocks:
            scaled_in_channels = make_divisible(in_channels * config.scale, config.divisor)
            scaled_out_channels = make_divisible(out_channels * config.scale, config.divisor)

            depthwise_block = PPLCNetDepthwiseSeparableConvLayer(
                in_channels=scaled_in_channels,
                out_channels=scaled_out_channels,
                kernel_size=kernel_size,
                stride=stride,
                use_squeeze_excitation=use_squeeze_excitation,
                config=config,
            )
            self.layers.append(depthwise_block)

    def forward(self, hidden_states):
        for layer in self.layers:
            hidden_states = layer(hidden_states)
        return hidden_states


@auto_docstring
class PPLCNetPreTrainedModel(PreTrainedModel):
    """
    An abstract base class for PP-LCNet models that inherits from Hugging Face PreTrainedModel.
    Provides common functionality for weight initialization and loading.
    """

    config: PPLCNetConfig
    base_model_prefix = "model"
    main_input_name = "pixel_values"
    input_modalities = ("image",)
    _can_compile_fullgraph = True
    supports_gradient_checkpointing = True
    _no_split_modules = ["PPLCNetDepthwiseSeparableConvLayer"]
    _can_record_outputs = {
        "hidden_states": PPLCNetBlock,
    }


class PPLCNetEncoder(PPLCNetPreTrainedModel):
    def __init__(self, config: PPLCNetConfig):
        super().__init__(config)
        self.config = config

        # stem
        self.convolution = PPLCNetConvLayer(
            in_channels=3,
            kernel_size=3,
            out_channels=make_divisible(config.stem_channels * config.scale, config.divisor),
            stride=config.stem_stride,
            activation=config.hidden_act,
        )
        # stages
        self.blocks = nn.ModuleList([])
        for stage_index in range(len(config.block_configs)):
            block = PPLCNetBlock(config, stage_index)
            self.blocks.append(block)

        self.post_init()

    @merge_with_config_defaults
    @capture_outputs
    def forward(self, pixel_values: torch.Tensor, **kwargs) -> BaseModelOutputWithNoAttention:
        hidden_state = self.convolution(pixel_values)
        for block in self.blocks:
            hidden_state = block(hidden_state)

        return BaseModelOutputWithNoAttention(last_hidden_state=hidden_state)


@auto_docstring(
    custom_intro="""
    PPLCNet backbone model for feature extraction.
    """
)
class PPLCNetBackbone(BackboneMixin, PPLCNetPreTrainedModel):
    has_attentions = False

    def __init__(self, config: PPLCNetConfig):
        super().__init__(config)
        num_features = [config.stem_channels]
        for block in config.block_configs:
            num_features.append(int(block[-1][2] * config.scale))
        self.num_features = num_features
        self.encoder = PPLCNetEncoder(config)

        self.post_init()

    @can_return_tuple
    @filter_output_hidden_states
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.Tensor,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BackboneOutput:
        r"""
        Examples:

        ```python
        >>> from transformers import PPLCNetConfig, PPLCNetBackbone
        >>> import torch

        >>> config = PPLCNetConfig()
        >>> model = PPLCNetBackbone(config)

        >>> pixel_values = torch.randn(1, 3, 224, 224)

        >>> with torch.no_grad():
        ...     outputs = model(pixel_values)

        >>> feature_maps = outputs.feature_maps
        >>> list(feature_maps[-1].shape)
        ```"""
        kwargs["output_hidden_states"] = True  # required to extract layers for the stages
        hidden_states = self.encoder(pixel_values, **kwargs).hidden_states

        feature_maps = ()
        for idx, stage in enumerate(self.stage_names):
            if stage in self.out_features:
                feature_maps += (hidden_states[idx],)

        return BackboneOutput(
            feature_maps=feature_maps,
            hidden_states=hidden_states,
        )


@auto_docstring
class PPLCNetForImageClassification(PPLCNetPreTrainedModel):
    _keys_to_ignore_on_load_missing = ["num_batches_tracked"]

    def __init__(self, config: PPLCNetConfig):
        super().__init__(config)
        self.encoder = PPLCNetEncoder(config)

        self.config = config
        self.num_labels = config.num_labels
        last_block_out_channels = config.block_configs[-1][-1][2]
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.last_convolution = nn.Conv2d(
            in_channels=make_divisible(last_block_out_channels * config.scale, config.divisor),
            out_channels=config.class_expand,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=False,
        )
        self.act_fn = ACT2FN[config.hidden_act]
        self.hidden_dropout_prob = config.hidden_dropout_prob

        self.flatten = nn.Flatten(start_dim=1, end_dim=-1)
        self.head = nn.Linear(config.class_expand, config.num_labels) if config.num_labels > 0 else nn.Identity()

        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithNoAttention:
        r"""
        Examples:

        ```python
        >>> import requests
        >>> from PIL import Image
        >>> from transformers import AutoModelForImageClassification, AutoImageProcessor

        >>> model_path = "PaddlePaddle/PP-LCNet_x1_0_table_cls_safetensors"
        >>> model = AutoModelForImageClassification.from_pretrained(model_path)
        >>> image_processor = AutoImageProcessor.from_pretrained(model_path)

        >>> url = "https://paddle-model-ecology.bj.bcebos.com/paddlex/imgs/demo_image/img_rot180_demo.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> inputs = image_processor(images=image, return_tensors="pt")
        >>> outputs = model(**inputs)
        >>> predicted_label = outputs.last_hidden_state.argmax(-1).item()
        >>> print(model.config.id2label[predicted_label])
        wireless_table
        ```"""
        outputs = self.encoder(pixel_values, **kwargs)

        last_hidden_state = self.avg_pool(outputs.last_hidden_state)

        last_hidden_state = self.last_convolution(last_hidden_state)
        last_hidden_state = self.act_fn(last_hidden_state)
        last_hidden_state = last_hidden_state * (1 - self.hidden_dropout_prob)

        last_hidden_state = self.flatten(last_hidden_state)
        last_hidden_state = self.head(last_hidden_state)

        return BaseModelOutputWithNoAttention(last_hidden_state=last_hidden_state, hidden_states=outputs.hidden_states)


__all__ = ["PPLCNetBackbone", "PPLCNetForImageClassification", "PPLCNetPreTrainedModel"]
