






'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetImgaImgPipeline,
    KandinskyVaaPriorEmbaEmbPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()
class 				lowercase       (							_lowerCamelCase ,       unittest.TestCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							KandinskyVaaControlnetImgaImgPipeline
    lowerCAmelCase__      				=							["image_embeds", "negative_image_embeds", "image", "hint"]
    lowerCAmelCase__      				=							["image_embeds", "negative_image_embeds", "image", "hint"]
    lowerCAmelCase__      				=							[
        "generator",
        "height",
        "width",
        "strength",
        "guidance_scale",
        "num_inference_steps",
        "return_dict",
        "guidance_scale",
        "num_images_per_prompt",
        "output_type",
        "return_dict",
    ]
    lowerCAmelCase__      				=							False

    @property
    def    a__				(					self     :     int     )							->							Optional[int]:

        '''simple docstring'''




        return 32

    @property
    def    a__				(					self     :     Any     )							->							List[Any]:

        '''simple docstring'''




        return 32

    @property
    def    a__				(					self     :     Tuple     )							->							int:

        '''simple docstring'''




        return self.time_input_dim

    @property
    def    a__				(					self     :     List[Any]     )							->							Union[str, Any]:

        '''simple docstring'''




        return self.time_input_dim * 4

    @property
    def    a__				(					self     :     List[Any]     )							->							Dict:

        '''simple docstring'''




        return 100

    @property
    def    a__				(					self     :     List[str]     )							->							List[Any]:

        '''simple docstring'''




        torch.manual_seed(0     )

        lowerCamelCase__  =     {
            "in_channels": 8,
            # Out channels is double in channels because predicts mean and variance
            "out_channels": 8,
            "addition_embed_type": "image_hint",
            "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
            "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
            "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "layers_per_block": 1,
            "encoder_hid_dim": self.text_embedder_hidden_size,
            "encoder_hid_dim_type": "image_proj",
            "cross_attention_dim": self.cross_attention_dim,
            "attention_head_dim": 4,
            "resnet_time_scale_shift": "scale_shift",
            "class_embed_type": None,
        }

        lowerCamelCase__  =     UNetaDConditionModel(**__lowerCamelCase     )
        return model

    @property
    def    a__				(					self     :     List[Any]     )							->							Dict:

        '''simple docstring'''




        return {
            "block_out_channels": [32, 32, 64, 64],
            "down_block_types": [
                "DownEncoderBlock2D",
                "DownEncoderBlock2D",
                "DownEncoderBlock2D",
                "AttnDownEncoderBlock2D",
            ],
            "in_channels": 3,
            "latent_channels": 4,
            "layers_per_block": 1,
            "norm_num_groups": 8,
            "norm_type": "spatial",
            "num_vq_embeddings": 12,
            "out_channels": 3,
            "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
            "vq_embed_dim": 4,
        }

    @property
    def    a__				(					self     :     Dict     )							->							str:

        '''simple docstring'''




        torch.manual_seed(0     )
        lowerCamelCase__  =     VQModel(**self.dummy_movq_kwargs     )
        return model

    def    a__				(					self     :     Dict     )							->							Dict:

        '''simple docstring'''




        lowerCamelCase__  =     self.dummy_unet
        lowerCamelCase__  =     self.dummy_movq

        lowerCamelCase__  =     {
            "num_train_timesteps": 1000,
            "beta_schedule": "linear",
            "beta_start": 0.0_0_0_8_5,
            "beta_end": 0.0_1_2,
            "clip_sample": False,
            "set_alpha_to_one": False,
            "steps_offset": 0,
            "prediction_type": "epsilon",
            "thresholding": False,
        }

        lowerCamelCase__  =     DDIMScheduler(**__lowerCamelCase     )

        lowerCamelCase__  =     {
            "unet": unet,
            "scheduler": scheduler,
            "movq": movq,
        }

        return components

    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     List[Any]=0     )							->							Optional[Any]:

        '''simple docstring'''




        lowerCamelCase__  =     floats_tensor((1, self.text_embedder_hidden_size)				,     rng=random.Random(__lowerCamelCase     )     ).to(__lowerCamelCase     )
        lowerCamelCase__  =     floats_tensor((1, self.text_embedder_hidden_size)				,     rng=random.Random(seed + 1     )     ).to(
            __lowerCamelCase     )
        # create init_image
        lowerCamelCase__  =     floats_tensor((1, 3, 64, 64)				,     rng=random.Random(__lowerCamelCase     )     ).to(__lowerCamelCase     )
        lowerCamelCase__  =     image.cpu().permute(0				,     2				,     3				,     1     )[0]
        lowerCamelCase__  =     Image.fromarray(np.uinta(__lowerCamelCase     )     ).convert("RGB"     ).resize((256, 256)     )
        # create hint
        lowerCamelCase__  =     floats_tensor((1, 3, 64, 64)				,     rng=random.Random(__lowerCamelCase     )     ).to(__lowerCamelCase     )

        if str(__lowerCamelCase     ).startswith("mps"     ):
            lowerCamelCase__  =     torch.manual_seed(__lowerCamelCase     )
        else:
            lowerCamelCase__  =     torch.Generator(device=__lowerCamelCase     ).manual_seed(__lowerCamelCase     )
        lowerCamelCase__  =     {
            "image": init_image,
            "image_embeds": image_embeds,
            "negative_image_embeds": negative_image_embeds,
            "hint": hint,
            "generator": generator,
            "height": 64,
            "width": 64,
            "num_inference_steps": 10,
            "guidance_scale": 7.0,
            "strength": 0.2,
            "output_type": "np",
        }
        return inputs




    def    a__				(					self     :     Dict     )							->							Optional[Any]:

        '''simple docstring'''




        lowerCamelCase__  =     "cpu"

        lowerCamelCase__  =     self.get_dummy_components()

        lowerCamelCase__  =     self.pipeline_class(**__lowerCamelCase     )
        lowerCamelCase__  =     pipe.to(__lowerCamelCase     )

        pipe.set_progress_bar_config(disable=__lowerCamelCase     )

        lowerCamelCase__  =     pipe(**self.get_dummy_inputs(__lowerCamelCase     )     )
        lowerCamelCase__  =     output.images

        lowerCamelCase__  =     pipe(
            **self.get_dummy_inputs(__lowerCamelCase     )				,     return_dict=__lowerCamelCase				,     )[0]

        lowerCamelCase__  =     image[0, -3:, -3:, -1]
        lowerCamelCase__  =     image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        lowerCamelCase__  =     np.array(
            [0.5_4_9_8_5_0_3_4, 0.5_5_5_0_9_3_6_5, 0.5_2_5_6_1_5_0_4, 0.5_5_7_0_4_9_4, 0.5_5_9_3_8_1_8, 0.5_2_6_3_9_7_9, 0.5_0_2_8_5_6_4_3, 0.5_0_6_9_8_4_6, 0.5_1_1_9_6_7_3_6]     )
        assert (
            np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2
        ), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
        assert (
            np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1E-2
        ), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''
@slow
@require_torch_gpu
class 				lowercase       (							unittest.TestCase      ):


    '''simple docstring'''

    def    a__				(					self     :     Any     )							->							List[Any]:

        '''simple docstring'''




        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




    def    a__				(					self     :     Tuple     )							->							Optional[Any]:

        '''simple docstring'''




        lowerCamelCase__  =     load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/kandinskyv22/kandinskyv22_controlnet_img2img_robotcat_fp16.npy"     )

        lowerCamelCase__  =     load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"     )
        lowerCamelCase__  =     init_image.resize((512, 512)     )

        lowerCamelCase__  =     load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/kandinskyv22/hint_image_cat.png"     )
        lowerCamelCase__  =     torch.from_numpy(np.array(__lowerCamelCase     )     ).float() / 255.0
        lowerCamelCase__  =     hint.permute(2				,     0				,     1     ).unsqueeze(0     )

        lowerCamelCase__  =     "A robot, 4k photo"

        lowerCamelCase__  =     KandinskyVaaPriorEmbaEmbPipeline.from_pretrained(
            "kandinsky-community/kandinsky-2-2-prior"				,     torch_dtype=torch.floataa     )
        pipe_prior.to(__lowerCamelCase     )

        lowerCamelCase__  =     KandinskyVaaControlnetImgaImgPipeline.from_pretrained(
            "kandinsky-community/kandinsky-2-2-controlnet-depth"				,     torch_dtype=torch.floataa     )
        lowerCamelCase__  =     pipeline.to(__lowerCamelCase     )

        pipeline.set_progress_bar_config(disable=__lowerCamelCase     )

        lowerCamelCase__  =     torch.Generator(device="cpu"     ).manual_seed(0     )

        lowerCamelCase__  ,	lowerCamelCase__  =     pipe_prior(
            __lowerCamelCase				,     image=__lowerCamelCase				,     strength=0.8_5				,     generator=__lowerCamelCase				,     negative_prompt=""				,     ).to_tuple()

        lowerCamelCase__  =     pipeline(
            image=__lowerCamelCase				,     image_embeds=__lowerCamelCase				,     negative_image_embeds=__lowerCamelCase				,     hint=__lowerCamelCase				,     generator=__lowerCamelCase				,     num_inference_steps=100				,     height=512				,     width=512				,     strength=0.5				,     output_type="np"				,     )

        lowerCamelCase__  =     output.images[0]

        assert image.shape == (512, 512, 3)

        assert_mean_pixel_difference(__lowerCamelCase				,     __lowerCamelCase     )









'''simple docstring'''





import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin
class 				lowercase       (							_lowerCamelCase ,       _lowerCamelCase      ):


    '''simple docstring'''

    @register_to_config
    def __init__(					self     :     Optional[Any]				,     __lowerCamelCase     :     int = 128				,     __lowerCamelCase     :     int = 256				,     __lowerCamelCase     :     float = 2_0_0_0.0				,     __lowerCamelCase     :     int = 768				,     __lowerCamelCase     :     int = 12				,     __lowerCamelCase     :     int = 12				,     __lowerCamelCase     :     int = 64				,     __lowerCamelCase     :     int = 2048				,     __lowerCamelCase     :     float = 0.1				,     )							->							str:

     '''simple docstring'''




     super().__init__()

     lowerCamelCase__  =     nn.Sequential(
         nn.Linear(__lowerCamelCase				,     d_model * 4				,     bias=__lowerCamelCase     )				,     nn.SiLU()				,     nn.Linear(d_model * 4				,     d_model * 4				,     bias=__lowerCamelCase     )				,     nn.SiLU()				,     )

     lowerCamelCase__  =     nn.Embedding(__lowerCamelCase				,     __lowerCamelCase     )
     lowerCamelCase__  =     False

     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )

     lowerCamelCase__  =     nn.Dropout(p=__lowerCamelCase     )

     lowerCamelCase__  =     nn.ModuleList()
     for lyr_num in range(__lowerCamelCase     ):
      # FiLM conditional T5 decoder
      lowerCamelCase__  =     DecoderLayer(d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )
      self.decoders.append(__lowerCamelCase     )

     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase     )

     lowerCamelCase__  =     nn.Dropout(p=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Optional[int]     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     torch.mul(query_input.unsqueeze(-1     )				,     key_input.unsqueeze(-2     )     )
     return mask.unsqueeze(-3     )




    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Union[str, Any]     )							->							Optional[int]:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     decoder_input_tokens.shape
     assert decoder_noise_time.shape == (batch,)

     # decoder_noise_time is in [0, 1), so rescale to expected timing range.
     lowerCamelCase__  =     get_timestep_embedding(
         decoder_noise_time * self.config.max_decoder_noise_time				,     embedding_dim=self.config.d_model				,     max_period=self.config.max_decoder_noise_time				,     ).to(dtype=self.dtype     )

     lowerCamelCase__  =     self.conditioning_emb(__lowerCamelCase     ).unsqueeze(1     )

     assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

     lowerCamelCase__  =     decoder_input_tokens.shape[1]

     # If we want to use relative positions for audio context, we can just offset
     # this sequence by the length of encodings_and_masks.
     lowerCamelCase__  =     torch.broadcast_to(
         torch.arange(__lowerCamelCase				,     device=decoder_input_tokens.device     )				,     (batch, seq_length)				,     )

     lowerCamelCase__  =     self.position_encoding(__lowerCamelCase     )

     lowerCamelCase__  =     self.continuous_inputs_projection(__lowerCamelCase     )
     inputs += position_encodings
     lowerCamelCase__  =     self.dropout(__lowerCamelCase     )

     # decoder: No padding present.
     lowerCamelCase__  =     torch.ones(
         decoder_input_tokens.shape[:2]				,     device=decoder_input_tokens.device				,     dtype=inputs.dtype     )

     # Translate encoding masks to encoder-decoder masks.
     lowerCamelCase__  =     [(x, self.encoder_decoder_mask(__lowerCamelCase				,     __lowerCamelCase     )) for x, y in encodings_and_masks]

     # cross attend style: concat encodings
     lowerCamelCase__  =     torch.cat([x[0] for x in encodings_and_encdec_masks]				,     dim=1     )
     lowerCamelCase__  =     torch.cat([x[1] for x in encodings_and_encdec_masks]				,     dim=-1     )

     for lyr in self.decoders:
      lowerCamelCase__  =     lyr(
          __lowerCamelCase				,     conditioning_emb=__lowerCamelCase				,     encoder_hidden_states=__lowerCamelCase				,     encoder_attention_mask=__lowerCamelCase				,     )[0]

     lowerCamelCase__  =     self.decoder_norm(__lowerCamelCase     )
     lowerCamelCase__  =     self.post_dropout(__lowerCamelCase     )

     lowerCamelCase__  =     self.spec_out(__lowerCamelCase     )
     return spec_out
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     str=1E-6     )							->							int:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.ModuleList()

     # cond self attention: layer 0
     self.layer.append(
         TaLayerSelfAttentionCond(d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )     )

     # cross attention: layer 1
     self.layer.append(
         TaLayerCrossAttention(
             d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     dropout_rate=__lowerCamelCase				,     layer_norm_epsilon=__lowerCamelCase				,     )     )

     # Film Cond MLP + dropout: last layer
     self.layer.append(
         TaLayerFFCond(d_model=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase				,     layer_norm_epsilon=__lowerCamelCase     )     )




    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     Optional[int]=None				,     __lowerCamelCase     :     int=None				,     __lowerCamelCase     :     List[Any]=None				,     __lowerCamelCase     :     List[str]=None				,     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer[0](
         __lowerCamelCase				,     conditioning_emb=__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     )

     if encoder_hidden_states is not None:
      lowerCamelCase__  =     torch.where(encoder_attention_mask > 0				,     0				,     -1E10     ).to(
          encoder_hidden_states.dtype     )

      lowerCamelCase__  =     self.layer[1](
          __lowerCamelCase				,     key_value_states=__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     )

     # Apply Film Conditional Feed Forward layer
     lowerCamelCase__  =     self.layer[-1](__lowerCamelCase				,     __lowerCamelCase     )

     return (hidden_states,)
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[Any]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Tuple     )							->							Optional[int]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase     )
     lowerCamelCase__  =     TaFiLMLayer(in_features=d_model * 4				,     out_features=__lowerCamelCase     )
     lowerCamelCase__  =     Attention(query_dim=__lowerCamelCase				,     heads=__lowerCamelCase				,     dim_head=__lowerCamelCase				,     out_bias=__lowerCamelCase				,     scale_qk=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     str=None				,     )							->							List[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )

     if conditioning_emb is not None:
      lowerCamelCase__  =     self.FiLMLayer(__lowerCamelCase				,     __lowerCamelCase     )

     # Self-attention block
     lowerCamelCase__  =     self.attention(__lowerCamelCase     )

     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )

     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Any				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Dict     )							->							Union[str, Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     Attention(query_dim=__lowerCamelCase				,     heads=__lowerCamelCase				,     dim_head=__lowerCamelCase				,     out_bias=__lowerCamelCase				,     scale_qk=__lowerCamelCase     )
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase				,     eps=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     List[str]				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     Dict=None				,     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )
     lowerCamelCase__  =     self.attention(
         __lowerCamelCase				,     encoder_hidden_states=__lowerCamelCase				,     attention_mask=attention_mask.squeeze(1     )				,     )
     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )
     return layer_output
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[str]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[str]     )							->							Optional[Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     TaDenseGatedActDense(d_model=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )
     lowerCamelCase__  =     TaFiLMLayer(in_features=d_model * 4				,     out_features=__lowerCamelCase     )
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase				,     eps=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     Any				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Union[str, Any]=None     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )
     if conditioning_emb is not None:
      lowerCamelCase__  =     self.film(__lowerCamelCase				,     __lowerCamelCase     )

     lowerCamelCase__  =     self.DenseReluDense(__lowerCamelCase     )
     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )
     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[str]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     str     )							->							str:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )
     lowerCamelCase__  =     NewGELUActivation()




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Union[str, Any]     )							->							Optional[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     self.act(self.wi_a(__lowerCamelCase     )     )
     lowerCamelCase__  =     self.wi_a(__lowerCamelCase     )
     lowerCamelCase__  =     hidden_gelu * hidden_linear
     lowerCamelCase__  =     self.dropout(__lowerCamelCase     )

     lowerCamelCase__  =     self.wo(__lowerCamelCase     )
     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     str				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     List[Any]=1E-6     )							->							Union[str, Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Parameter(torch.ones(__lowerCamelCase     )     )
     lowerCamelCase__  =     eps




    def    a__				(					self     :     Optional[int]				,     __lowerCamelCase     :     str     )							->							int:

     '''simple docstring'''




     lowerCamelCase__  =     hidden_states.to(torch.floataa     ).pow(2     ).mean(-1				,     keepdim=__lowerCamelCase     )
     lowerCamelCase__  =     hidden_states * torch.rsqrt(variance + self.variance_epsilon     )

     # convert into half-precision if necessary
     if self.weight.dtype in [torch.floataa, torch.bfloataa]:
      lowerCamelCase__  =     hidden_states.to(self.weight.dtype     )

     return self.weight * hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     torch.Tensor     )							->							torch.Tensor:

     '''simple docstring'''




     return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi     ) * (input + 0.0_4_4_7_1_5 * torch.pow(__lowerCamelCase				,     3.0     ))     ))
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Optional[int]				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]     )							->							Optional[int]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     out_features * 2				,     bias=__lowerCamelCase     )




    def    a__				(					self     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     str     )							->							List[str]:

     '''simple docstring'''




     lowerCamelCase__  =     self.scale_bias(__lowerCamelCase     )
     lowerCamelCase__  ,	lowerCamelCase__  =     torch.chunk(__lowerCamelCase				,     2				,     -1     )
     lowerCamelCase__  =     x * (1 + scale) + shift
     return x

