'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
from .upsampling import upfirdn2d_native
from .normalization import RMSNorm
from typing import Optional, Tuple
import torch.nn.functional as F
from ..utils import deprecate
import torch.nn as nn
import torch
class Downsample1D(nn.Module):
    def __init__(self, channels: int, use_conv: bool=False, out_channels: Optional[int]=None, padding: int=1, name: str='conv'):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_conv = use_conv
        self.padding = padding
        stride = 2
        self.name = name
        if use_conv: self.conv = nn.Conv1d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
        else:
            assert self.channels == self.out_channels
            self.conv = nn.AvgPool1d(kernel_size=stride, stride=stride)
    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        assert inputs.shape[1] == self.channels
        return self.conv(inputs)
class Downsample2D(nn.Module):
    def __init__(self, channels: int, use_conv: bool=False, out_channels: Optional[int]=None, padding: int=1, name: str='conv', kernel_size=3, norm_type=None, eps=None, elementwise_affine=None, bias=True):
        super().__init__()
        self.channels = channels
        self.out_channels = out_channels or channels
        self.use_conv = use_conv
        self.padding = padding
        stride = 2
        self.name = name
        if norm_type == 'ln_norm': self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
        elif norm_type == 'rms_norm': self.norm = RMSNorm(channels, eps, elementwise_affine)
        elif norm_type is None: self.norm = None
        else: raise ValueError(f'unknown norm_type: {norm_type}')
        if use_conv: conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
        else:
            assert self.channels == self.out_channels
            conv = nn.AvgPool2d(kernel_size=stride, stride=stride)
        if name == 'conv':
            self.Conv2d_0 = conv
            self.conv = conv
        elif name == 'Conv2d_0': self.conv = conv
        else: self.conv = conv
    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        assert hidden_states.shape[1] == self.channels
        if self.norm is not None: hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
        if self.use_conv and self.padding == 0:
            pad = (0, 1, 0, 1)
            hidden_states = F.pad(hidden_states, pad, mode='constant', value=0)
        assert hidden_states.shape[1] == self.channels
        hidden_states = self.conv(hidden_states)
        return hidden_states
class FirDownsample2D(nn.Module):
    def __init__(self, channels: Optional[int]=None, out_channels: Optional[int]=None, use_conv: bool=False, fir_kernel: Tuple[int, int, int, int]=(1, 3, 3, 1)):
        super().__init__()
        out_channels = out_channels if out_channels else channels
        if use_conv: self.Conv2d_0 = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.fir_kernel = fir_kernel
        self.use_conv = use_conv
        self.out_channels = out_channels
    def _downsample_2d(self, hidden_states: torch.Tensor, weight: Optional[torch.Tensor]=None, kernel: Optional[torch.Tensor]=None, factor: int=2, gain: float=1) -> torch.Tensor:
        """Returns:"""
        assert isinstance(factor, int) and factor >= 1
        if kernel is None: kernel = [1] * factor
        kernel = torch.tensor(kernel, dtype=torch.float32)
        if kernel.ndim == 1: kernel = torch.outer(kernel, kernel)
        kernel /= torch.sum(kernel)
        kernel = kernel * gain
        if self.use_conv:
            _, _, convH, convW = weight.shape
            pad_value = kernel.shape[0] - factor + (convW - 1)
            stride_value = [factor, factor]
            upfirdn_input = upfirdn2d_native(hidden_states, torch.tensor(kernel, device=hidden_states.device), pad=((pad_value + 1) // 2, pad_value // 2))
            output = F.conv2d(upfirdn_input, weight, stride=stride_value, padding=0)
        else:
            pad_value = kernel.shape[0] - factor
            output = upfirdn2d_native(hidden_states, torch.tensor(kernel, device=hidden_states.device), down=factor, pad=((pad_value + 1) // 2, pad_value // 2))
        return output
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        if self.use_conv:
            downsample_input = self._downsample_2d(hidden_states, weight=self.Conv2d_0.weight, kernel=self.fir_kernel)
            hidden_states = downsample_input + self.Conv2d_0.bias.reshape(1, -1, 1, 1)
        else: hidden_states = self._downsample_2d(hidden_states, kernel=self.fir_kernel, factor=2)
        return hidden_states
class KDownsample2D(nn.Module):
    def __init__(self, pad_mode: str='reflect'):
        super().__init__()
        self.pad_mode = pad_mode
        kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]])
        self.pad = kernel_1d.shape[1] // 2 - 1
        self.register_buffer('kernel', kernel_1d.T @ kernel_1d, persistent=False)
    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        inputs = F.pad(inputs, (self.pad,) * 4, self.pad_mode)
        weight = inputs.new_zeros([inputs.shape[1], inputs.shape[1], self.kernel.shape[0], self.kernel.shape[1]])
        indices = torch.arange(inputs.shape[1], device=inputs.device)
        kernel = self.kernel.to(weight)[None, :].expand(inputs.shape[1], -1, -1)
        weight[indices, indices] = kernel
        return F.conv2d(inputs, weight, stride=2)
class CogVideoXDownsample3D(nn.Module):
    """Args:"""
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int=3, stride: int=2, padding: int=0, compress_time: bool=False):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
        self.compress_time = compress_time
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.compress_time:
            batch_size, channels, frames, height, width = x.shape
            x = x.permute(0, 3, 4, 1, 2).reshape(batch_size * height * width, channels, frames)
            if x.shape[-1] % 2 == 1:
                x_first, x_rest = (x[..., 0], x[..., 1:])
                if x_rest.shape[-1] > 0: x_rest = F.avg_pool1d(x_rest, kernel_size=2, stride=2)
                x = torch.cat([x_first[..., None], x_rest], dim=-1)
                x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2)
            else:
                x = F.avg_pool1d(x, kernel_size=2, stride=2)
                x = x.reshape(batch_size, height, width, channels, x.shape[-1]).permute(0, 3, 4, 1, 2)
        pad = (0, 1, 0, 1)
        x = F.pad(x, pad, mode='constant', value=0)
        batch_size, channels, frames, height, width = x.shape
        x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channels, height, width)
        x = self.conv(x)
        x = x.reshape(batch_size, frames, x.shape[1], x.shape[2], x.shape[3]).permute(0, 2, 1, 3, 4)
        return x
def downsample_2d(hidden_states: torch.Tensor, kernel: Optional[torch.Tensor]=None, factor: int=2, gain: float=1) -> torch.Tensor:
    """Returns:"""
    assert isinstance(factor, int) and factor >= 1
    if kernel is None: kernel = [1] * factor
    kernel = torch.tensor(kernel, dtype=torch.float32)
    if kernel.ndim == 1: kernel = torch.outer(kernel, kernel)
    kernel /= torch.sum(kernel)
    kernel = kernel * gain
    pad_value = kernel.shape[0] - factor
    output = upfirdn2d_native(hidden_states, kernel.to(device=hidden_states.device), down=factor, pad=((pad_value + 1) // 2, pad_value // 2))
    return output
'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
