


from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


SCREAMING_SNAKE_CASE       :		Optional[Any]         =				logging.get_logger(__name__)

SCREAMING_SNAKE_CASE       :		List[str]         =				{
    "Helsinki-NLP/opus-mt-en-de": "https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json",
    # See all Marian models at https://huggingface.co/models?filter=marian
}






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   lowercase					:	Union[str, Any]									="""marian"""
   lowercase					:	List[Any]									=["""past_key_values"""]
   lowercase					:	int									={"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}
   def __init__(     self   ,							UpperCamelCase_=5_8101   ,							UpperCamelCase_=None   ,							UpperCamelCase_=1024   ,							UpperCamelCase_=12   ,							UpperCamelCase_=4096   ,							UpperCamelCase_=16   ,							UpperCamelCase_=12   ,							UpperCamelCase_=4096   ,							UpperCamelCase_=16   ,							UpperCamelCase_=0.0   ,							UpperCamelCase_=0.0   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_="gelu"   ,							UpperCamelCase_=1024   ,							UpperCamelCase_=0.1   ,							UpperCamelCase_=0.0   ,							UpperCamelCase_=0.0   ,							UpperCamelCase_=0.02   ,							UpperCamelCase_=5_8100   ,							UpperCamelCase_=False   ,							UpperCamelCase_=5_8100   ,							UpperCamelCase_=0   ,							UpperCamelCase_=0   ,							UpperCamelCase_=True   ,							**UpperCamelCase_   ,							):
        lowercase_       :Union[str, Any]     =       vocab_size
        lowercase_       :Tuple     =       decoder_vocab_size or vocab_size
        lowercase_       :List[Any]     =       max_position_embeddings
        lowercase_       :Any     =       d_model
        lowercase_       :Tuple     =       encoder_ffn_dim
        lowercase_       :Union[str, Any]     =       encoder_layers
        lowercase_       :Tuple     =       encoder_attention_heads
        lowercase_       :List[str]     =       decoder_ffn_dim
        lowercase_       :int     =       decoder_layers
        lowercase_       :List[Any]     =       decoder_attention_heads
        lowercase_       :Any     =       dropout
        lowercase_       :Union[str, Any]     =       attention_dropout
        lowercase_       :Any     =       activation_dropout
        lowercase_       :Union[str, Any]     =       activation_function
        lowercase_       :List[Any]     =       init_std
        lowercase_       :List[Any]     =       encoder_layerdrop
        lowercase_       :Any     =       decoder_layerdrop
        lowercase_       :List[Any]     =       use_cache
        lowercase_       :List[Any]     =       encoder_layers
        lowercase_       :Any     =       scale_embedding  # scale factor will be sqrt(d_model) if True
        lowercase_       :List[str]     =       share_encoder_decoder_embeddings
        super().__init__(
            pad_token_id=UpperCamelCase_   ,							eos_token_id=UpperCamelCase_   ,							is_encoder_decoder=UpperCamelCase_   ,							decoder_start_token_id=UpperCamelCase_   ,							forced_eos_token_id=UpperCamelCase_   ,							**UpperCamelCase_   ,							)






class      UpperCamelCase					(		lowercase__     ):


   '''simple docstring'''







   @property
   # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
   def 		UpperCamelCase  (     self		):
        if self.task in ["default", "seq2seq-lm"]:
             lowercase_       :int     =       OrderedDict(
                 [
                     ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                     ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                 ]		)

             if self.use_past:
                  lowercase_       :Optional[int]     =       {0: '''batch'''}
                  lowercase_       :Optional[Any]     =       {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
             else:
                  lowercase_       :Tuple     =       {0: '''batch''', 1: '''decoder_sequence'''}
                  lowercase_       :Dict     =       {0: '''batch''', 1: '''decoder_sequence'''}

             if self.use_past:
                  self.fill_with_past_key_values_(UpperCamelCase_   ,							direction='''inputs'''		)
        elif self.task == "causal-lm":
             # TODO: figure this case out.
             lowercase_       :int     =       OrderedDict(
                 [
                     ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                     ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                 ]		)
             if self.use_past:
                  lowercase_    ,      lowercase_       :Union[str, Any]     =       self.num_layers
                  for i in range(UpperCamelCase_		):
                       lowercase_       :Dict     =       {0: '''batch''', 2: '''past_sequence + sequence'''}
                       lowercase_       :Tuple     =       {0: '''batch''', 2: '''past_sequence + sequence'''}
        else:
             lowercase_       :str     =       OrderedDict(
                 [
                     ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                     ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                     ('''decoder_input_ids''', {0: '''batch''', 1: '''decoder_sequence'''}),
                     ('''decoder_attention_mask''', {0: '''batch''', 1: '''decoder_sequence'''}),
                 ]		)

        return common_inputs
   @property
   # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
   def 		UpperCamelCase  (     self		):
        if self.task in ["default", "seq2seq-lm"]:
             lowercase_       :List[str]     =       super().outputs
        else:
             lowercase_       :str     =       super(UpperCamelCase_   ,							self		).outputs
             if self.use_past:
                  lowercase_    ,      lowercase_       :Dict     =       self.num_layers
                  for i in range(UpperCamelCase_		):
                       lowercase_       :str     =       {0: '''batch''', 2: '''past_sequence + sequence'''}
                       lowercase_       :str     =       {0: '''batch''', 2: '''past_sequence + sequence'''}
        return common_outputs
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = False   ,							UpperCamelCase_ = None   ,							):
        lowercase_       :Optional[Any]     =       self._generate_dummy_inputs_for_encoder_and_decoder(
            UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)

        # Generate decoder inputs
        lowercase_       :List[str]     =       seq_length if not self.use_past else 1
        lowercase_       :Optional[int]     =       self._generate_dummy_inputs_for_encoder_and_decoder(
            UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)
        lowercase_       :Union[str, Any]     =       {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
        lowercase_       :List[Any]     =       dict(**UpperCamelCase_   ,							**UpperCamelCase_		)

        if self.use_past:
             if not is_torch_available():
                  raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''		)
             else:
                  import torch
             lowercase_    ,      lowercase_       :Union[str, Any]     =       common_inputs['''input_ids'''].shape
             lowercase_       :Optional[Any]     =       common_inputs['''decoder_input_ids'''].shape[1]
             lowercase_    ,      lowercase_       :int     =       self.num_attention_heads
             lowercase_       :List[Any]     =       (
                 batch,
                 num_encoder_attention_heads,
                 encoder_seq_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             lowercase_       :int     =       decoder_seq_length + 3
             lowercase_       :str     =       (
                 batch,
                 num_decoder_attention_heads,
                 decoder_past_length,
                 self._config.hidden_size // num_decoder_attention_heads,
             )

             lowercase_       :Dict     =       torch.cat(
                 [common_inputs['''decoder_attention_mask'''], torch.ones(UpperCamelCase_   ,							UpperCamelCase_		)]   ,							dim=1		)

             lowercase_       :Optional[int]     =       []
             # If the number of encoder and decoder layers are present in the model configuration, both are considered
             lowercase_    ,      lowercase_       :Optional[int]     =       self.num_layers
             lowercase_       :Optional[int]     =       min(UpperCamelCase_   ,							UpperCamelCase_		)
             lowercase_       :Any     =       max(UpperCamelCase_   ,							UpperCamelCase_		) - min_num_layers
             lowercase_       :Optional[Any]     =       '''encoder''' if num_encoder_layers > num_decoder_layers else '''decoder'''

             for _ in range(UpperCamelCase_		):
                  common_inputs["past_key_values"].append(
                      (
                          torch.zeros(UpperCamelCase_		),
                          torch.zeros(UpperCamelCase_		),
                          torch.zeros(UpperCamelCase_		),
                          torch.zeros(UpperCamelCase_		),
                      )		)
             # TODO: test this.
             lowercase_       :Optional[Any]     =       encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
             for _ in range(UpperCamelCase_   ,							UpperCamelCase_		):
                  common_inputs["past_key_values"].append((torch.zeros(UpperCamelCase_		), torch.zeros(UpperCamelCase_		))		)
        return common_inputs
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = False   ,							UpperCamelCase_ = None   ,							):
        lowercase_       :Dict     =       self._generate_dummy_inputs_for_encoder_and_decoder(
            UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)

        if self.use_past:
             if not is_torch_available():
                  raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''		)
             else:
                  import torch
             lowercase_    ,      lowercase_       :Optional[int]     =       common_inputs['''input_ids'''].shape
             # Not using the same length for past_key_values
             lowercase_       :Union[str, Any]     =       seqlen + 2
             lowercase_    ,      lowercase_       :Tuple     =       self.num_layers
             lowercase_    ,      lowercase_       :Dict     =       self.num_attention_heads
             lowercase_       :List[str]     =       (
                 batch,
                 num_encoder_attention_heads,
                 past_key_values_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )

             lowercase_       :Tuple     =       common_inputs['''attention_mask'''].dtype
             lowercase_       :Tuple     =       torch.cat(
                 [common_inputs['''attention_mask'''], torch.ones(UpperCamelCase_   ,							UpperCamelCase_   ,							dtype=UpperCamelCase_		)]   ,							dim=1		)
             lowercase_       :Tuple     =       [
                 (torch.zeros(UpperCamelCase_		), torch.zeros(UpperCamelCase_		)) for _ in range(UpperCamelCase_		)
             ]
        return common_inputs
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = False   ,							UpperCamelCase_ = None   ,							):
        # Copied from OnnxConfig.generate_dummy_inputs
        # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
        # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
        lowercase_       :Dict     =       compute_effective_axis_dimension(
            UpperCamelCase_   ,							fixed_dimension=OnnxConfig.default_fixed_batch   ,							num_token_to_add=0		)

        # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
        lowercase_       :Any     =       tokenizer.num_special_tokens_to_add(UpperCamelCase_		)
        lowercase_       :Optional[Any]     =       compute_effective_axis_dimension(
            UpperCamelCase_   ,							fixed_dimension=OnnxConfig.default_fixed_sequence   ,							num_token_to_add=UpperCamelCase_		)

        # Generate dummy inputs according to compute batch and sequence
        lowercase_       :Optional[int]     =       [''' '''.join([tokenizer.unk_token]		) * seq_length] * batch_size
        lowercase_       :int     =       dict(tokenizer(UpperCamelCase_   ,							return_tensors=UpperCamelCase_		)		)
        return common_inputs
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = -1   ,							UpperCamelCase_ = False   ,							UpperCamelCase_ = None   ,							):
        if self.task in ["default", "seq2seq-lm"]:
             lowercase_       :Any     =       self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                 UpperCamelCase_   ,							batch_size=UpperCamelCase_   ,							seq_length=UpperCamelCase_   ,							is_pair=UpperCamelCase_   ,							framework=UpperCamelCase_		)

        else:
             lowercase_       :Dict     =       self._generate_dummy_inputs_for_causal_lm(
                 UpperCamelCase_   ,							batch_size=UpperCamelCase_   ,							seq_length=UpperCamelCase_   ,							is_pair=UpperCamelCase_   ,							framework=UpperCamelCase_		)

        return common_inputs
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        if self.task in ["default", "seq2seq-lm"]:
             lowercase_       :int     =       super()._flatten_past_key_values_(UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)
        else:
             lowercase_       :str     =       super(UpperCamelCase_   ,							self		)._flatten_past_key_values_(
                 UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)




   @property
   def 		UpperCamelCase  (     self		):
        return 1E-4





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


SCREAMING_SNAKE_CASE       :		int         =				False






class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   @property
   def 		UpperCamelCase  (     self		):
        return 12
   @property
   def 		UpperCamelCase  (     self		):
        return 12
   @property
   def 		UpperCamelCase  (     self		):
        return 32
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :List[Any]     =       VQModel(
            block_out_channels=[32, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,							up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,							latent_channels=3   ,							num_vq_embeddings=self.num_embed   ,							vq_embed_dim=3   ,							)
        return model
   @property
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)
        return tokenizer
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)
        lowercase_       :int     =       CLIPTextConfig(
            bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=self.text_embedder_hidden_size   ,							intermediate_size=37   ,							layer_norm_eps=1E-05   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=1000   ,							)
        return CLIPTextModel(UpperCamelCase_		)
   @property
   def 		UpperCamelCase  (     self		):
        torch.manual_seed(0		)

        lowercase_       :int     =       12
        lowercase_       :List[Any]     =       12

        lowercase_       :Dict     =       {
            '''attention_bias''': True,
            '''cross_attention_dim''': 32,
            '''attention_head_dim''': height * width,
            '''num_attention_heads''': 1,
            '''num_vector_embeds''': self.num_embed,
            '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
            '''norm_num_groups''': 32,
            '''sample_size''': width,
            '''activation_fn''': '''geglu-approximate''',
        }

        lowercase_       :int     =       TransformeraDModel(**UpperCamelCase_		)
        return model
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       '''cpu'''

        lowercase_       :int     =       self.dummy_vqvae
        lowercase_       :int     =       self.dummy_text_encoder
        lowercase_       :Any     =       self.dummy_tokenizer
        lowercase_       :Optional[int]     =       self.dummy_transformer
        lowercase_       :List[str]     =       VQDiffusionScheduler(self.num_embed		)
        lowercase_       :int     =       LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCamelCase_		)

        lowercase_       :List[Any]     =       VQDiffusionPipeline(
            vqvae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							transformer=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							learned_classifier_free_sampling_embeddings=UpperCamelCase_   ,							)
        lowercase_       :Union[str, Any]     =       pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :Dict     =       '''teddy bear playing in the pool'''

        lowercase_       :Dict     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       pipe([prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''		)
        lowercase_       :Any     =       output.images

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :str     =       pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							return_dict=UpperCamelCase_   ,							num_inference_steps=2		)[0]

        lowercase_       :Optional[Any]     =       image[0, -3:, -3:, -1]
        lowercase_       :Tuple     =       image_from_tuple[0, -3:, -3:, -1]

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

        lowercase_       :str     =       np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2




   def 		UpperCamelCase  (     self		):
        lowercase_       :int     =       '''cpu'''

        lowercase_       :Dict     =       self.dummy_vqvae
        lowercase_       :str     =       self.dummy_text_encoder
        lowercase_       :List[Any]     =       self.dummy_tokenizer
        lowercase_       :Any     =       self.dummy_transformer
        lowercase_       :Optional[Any]     =       VQDiffusionScheduler(self.num_embed		)
        lowercase_       :List[str]     =       LearnedClassifierFreeSamplingEmbeddings(
            learnable=UpperCamelCase_   ,							hidden_size=self.text_embedder_hidden_size   ,							length=tokenizer.model_max_length		)

        lowercase_       :Optional[int]     =       VQDiffusionPipeline(
            vqvae=UpperCamelCase_   ,							text_encoder=UpperCamelCase_   ,							tokenizer=UpperCamelCase_   ,							transformer=UpperCamelCase_   ,							scheduler=UpperCamelCase_   ,							learned_classifier_free_sampling_embeddings=UpperCamelCase_   ,							)
        lowercase_       :Dict     =       pipe.to(UpperCamelCase_		)
        pipe.set_progress_bar_config(disable=UpperCamelCase_		)

        lowercase_       :int     =       '''teddy bear playing in the pool'''

        lowercase_       :Tuple     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Optional[int]     =       pipe([prompt]   ,							generator=UpperCamelCase_   ,							num_inference_steps=2   ,							output_type='''np'''		)
        lowercase_       :Optional[Any]     =       output.images

        lowercase_       :Optional[int]     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :Dict     =       pipe(
            [prompt]   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							return_dict=UpperCamelCase_   ,							num_inference_steps=2		)[0]

        lowercase_       :List[str]     =       image[0, -3:, -3:, -1]
        lowercase_       :Optional[int]     =       image_from_tuple[0, -3:, -3:, -1]

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

        lowercase_       :Dict     =       np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]		)

        assert np.abs(image_slice.flatten() - expected_slice		).max() < 2.0
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2






@slow
@require_torch_gpu
class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self		):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




   def 		UpperCamelCase  (     self		):
        lowercase_       :List[str]     =       load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''		)

        lowercase_       :int     =       VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''		)
        lowercase_       :Tuple     =       pipeline.to(UpperCamelCase_		)
        pipeline.set_progress_bar_config(disable=UpperCamelCase_		)

        # requires GPU generator for gumbel softmax
        # don't use GPU generator in tests though
        lowercase_       :Dict     =       torch.Generator(device=UpperCamelCase_		).manual_seed(0		)
        lowercase_       :int     =       pipeline(
            '''teddy bear playing in the pool'''   ,							num_images_per_prompt=1   ,							generator=UpperCamelCase_   ,							output_type='''np'''   ,							)

        lowercase_       :List[str]     =       output.images[0]

        assert image.shape == (256, 256, 3)
        assert np.abs(expected_image - image		).max() < 2.0

