
"""simple docstring"""
from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


lowercase_       :						Dict     					=	logging.get_logger(__name__)


lowercase_       :						List[str]     					=	{
    '''Salesforce/codegen-350M-nl''': '''https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json''',
    '''Salesforce/codegen-350M-multi''': '''https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json''',
    '''Salesforce/codegen-350M-mono''': '''https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json''',
    '''Salesforce/codegen-2B-nl''': '''https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-2B-multi''': '''https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-2B-mono''': '''https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-6B-nl''': '''https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-6B-multi''': '''https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-6B-mono''': '''https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-16B-nl''': '''https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-16B-multi''': '''https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-16B-mono''': '''https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json''',
}
class 			UpperCamelCase      (							__a  ):
     A__      =			"""codegen"""
     A__      =			{
         """max_position_embeddings""": """n_positions""",
         """hidden_size""": """n_embd""",
         """num_attention_heads""": """n_head""",
         """num_hidden_layers""": """n_layer""",
     }
     def __init__(     self  ,     snake_case__=50400  ,     snake_case__=2048  ,     snake_case__=2048  ,     snake_case__=4096  ,     snake_case__=28  ,     snake_case__=16  ,     snake_case__=64  ,     snake_case__=None  ,     snake_case__="gelu_new"  ,     snake_case__=0.0  ,     snake_case__=0.0  ,     snake_case__=0.0  ,     snake_case__=1E-5  ,     snake_case__=0.02  ,     snake_case__=True  ,     snake_case__=50256  ,     snake_case__=50256  ,     snake_case__=False  ,     **snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = vocab_size
       _SCREAMING_SNAKE_CASE   :				Dict							    = n_ctx
       _SCREAMING_SNAKE_CASE   :				Dict							    = n_positions
       _SCREAMING_SNAKE_CASE   :				Dict							    = n_embd
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = n_layer
       _SCREAMING_SNAKE_CASE   :				List[str]							    = n_head
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = n_inner
       _SCREAMING_SNAKE_CASE   :				Dict							    = rotary_dim
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = activation_function
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = resid_pdrop
       _SCREAMING_SNAKE_CASE   :				Dict							    = embd_pdrop
       _SCREAMING_SNAKE_CASE   :				Dict							    = attn_pdrop
       _SCREAMING_SNAKE_CASE   :				str							    = layer_norm_epsilon
       _SCREAMING_SNAKE_CASE   :				int							    = initializer_range
       _SCREAMING_SNAKE_CASE   :				str							    = use_cache

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = bos_token_id
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = eos_token_id

       super().__init__(
           bos_token_id=A__  ,     eos_token_id=A__  ,     tie_word_embeddings=A__  ,     **A__				)
class 			UpperCamelCase      (							__a  ):
     def __init__(     self  ,     snake_case__  ,     snake_case__ = "default"  ,     snake_case__ = None  ,     snake_case__ = False  ,     ):




       """simple docstring"""


       super().__init__(A__  ,     task=A__  ,     patching_specs=A__  ,     use_past=A__				)
       if not getattr(self._config  ,     "pad_token_id"  ,     A__				):
         # TODO: how to do that better?
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = 0
     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}				)
       if self.use_past:
         self.fill_with_past_key_values_(A__  ,     direction="inputs"				)
         _SCREAMING_SNAKE_CASE   :				Any							    = {0: "batch", 1: "past_sequence + sequence"}
       else:
         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = {0: "batch", 1: "sequence"}

       return common_inputs
     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return self._config.n_layer
     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return self._config.n_head
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = -1  ,     snake_case__ = -1  ,     snake_case__ = False  ,     snake_case__ = None  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[Any]							    = super(A__  ,     self				).generate_dummy_inputs(
           A__  ,     batch_size=A__  ,     seq_length=A__  ,     is_pair=A__  ,     framework=A__				)

       # We need to order the input in the way they appears in the forward()
       _SCREAMING_SNAKE_CASE   :				int							    = OrderedDict({"input_ids": common_inputs["input_ids"]}				)

       # Need to add the past_keys
       if self.use_past:
         if not is_torch_available():
           raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."				)
         else:
           import torch

           _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				str							    = common_inputs["input_ids"].shape
           # Not using the same length for past_key_values
           _SCREAMING_SNAKE_CASE   :				List[str]							    = seqlen + 2
           _SCREAMING_SNAKE_CASE   :				str							    = (
               batch,
               self.num_attention_heads,
               past_key_values_length,
               self._config.hidden_size // self.num_attention_heads,
           )
           _SCREAMING_SNAKE_CASE   :				List[str]							    = [
               (torch.zeros(A__				), torch.zeros(A__				)) for _ in range(self.num_layers				)
           ]

       _SCREAMING_SNAKE_CASE   :				str							    = common_inputs["attention_mask"]
       if self.use_past:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = ordered_inputs["attention_mask"].dtype
         _SCREAMING_SNAKE_CASE   :				int							    = torch.cat(
             [ordered_inputs["attention_mask"], torch.ones(A__  ,     A__  ,     dtype=A__				)]  ,     dim=1				)

       return ordered_inputs




     @property
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return 13







"""simple docstring"""
import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
     A__      =			DiTPipeline
     A__      =			CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
     A__      =			PipelineTesterMixin.required_optional_params - {
         """latents""",
         """num_images_per_prompt""",
         """callback""",
         """callback_steps""",
     }
     A__      =			CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
     A__      =			False
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = TransformeraDModel(
           sample_size=16  ,     num_layers=2  ,     patch_size=4  ,     attention_head_dim=8  ,     num_attention_heads=2  ,     in_channels=4  ,     out_channels=8  ,     attention_bias=snake_case__  ,     activation_fn="gelu-approximate"  ,     num_embeds_ada_norm=1000  ,     norm_type="ada_norm_zero"  ,     norm_elementwise_affine=snake_case__  ,     )
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = AutoencoderKL()
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = DDIMScheduler()
       _SCREAMING_SNAKE_CASE   :				Dict							    = {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
       return components
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__=0				):




       """simple docstring"""


       if str(snake_case__				).startswith("mps"				):
         _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.manual_seed(snake_case__				)
       else:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = torch.Generator(device=snake_case__				).manual_seed(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = {
           "class_labels": [1],
           "generator": generator,
           "num_inference_steps": 2,
           "output_type": "numpy",
       }
       return inputs
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = "cpu"

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.get_dummy_components()
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.pipeline_class(**snake_case__				)
       pipe.to(snake_case__				)
       pipe.set_progress_bar_config(disable=snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.get_dummy_inputs(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = pipe(**snake_case__				).images
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = image[0, -3:, -3:, -1]

       self.assertEqual(image.shape  ,     (1, 16, 16, 3)				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = np.abs(image_slice.flatten() - expected_slice				).max()
       self.assertLessEqual(snake_case__  ,     1E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self._test_inference_batch_single_identical(relax_max_difference=snake_case__  ,     expected_max_diff=1E-3				)




     @unittest.skipIf(
         torch_device != "cuda" or not is_xformers_available()  ,     reason="XFormers attention is only available with CUDA and `xformers` installed"  ,     )
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3				)
@require_torch_gpu
@slow
class 			UpperCamelCase      (							unittest.TestCase  ):
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = torch.manual_seed(0				)

       _SCREAMING_SNAKE_CASE   :				str							    = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"				)
       pipe.to("cuda"				)

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = ["vase", "umbrella", "white shark", "white wolf"]
       _SCREAMING_SNAKE_CASE   :				Any							    = pipe.get_label_ids(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				List[str]							    = pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=40  ,     output_type="np"				).images

       for word, image in zip(snake_case__  ,     snake_case__				):
         _SCREAMING_SNAKE_CASE   :				int							    = load_numpy(
             F'''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'''				)
         assert np.abs((expected_image - image).max()				) < 1E-2




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"				)
       _SCREAMING_SNAKE_CASE   :				Dict							    = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
       pipe.to("cuda"				)

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = ["vase", "umbrella"]
       _SCREAMING_SNAKE_CASE   :				int							    = pipe.get_label_ids(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				str							    = torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=25  ,     output_type="np"				).images

       for word, image in zip(snake_case__  ,     snake_case__				):
         _SCREAMING_SNAKE_CASE   :				Tuple							    = load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             F'''/dit/{word}_512.npy'''				)

         assert np.abs((expected_image - image).max()				) < 1E-1



