
"""simple docstring"""



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
 import os

 import jax.numpy as jnp
 from jax import jit

 from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
 from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

 a					:	Optional[int]							=		'''0.12'''  # assumed parallelism: 8


if is_torch_available():
 import torch


def  _UpperCamelCase	(						_A					,      _A					,      _A=None      )    ->  Union[str, Any]:





    """simple docstring"""

    if rng is None:
        _UpperCAmelCase          =   random.Random()

    _UpperCAmelCase          =   1
    for dim in shape:
        total_dims *= dim

    _UpperCAmelCase          =   []
    for _ in range(_A      ):
        values.append(rng.randint(0					,      vocab_size - 1      )      )

    _UpperCAmelCase          =   np.array(_A					,      dtype=jnp.intaa      ).reshape(_A      )

    return output


def  _UpperCamelCase	(						_A					,      _A=None      )    ->  Tuple:





    """simple docstring"""

    _UpperCAmelCase          =   ids_tensor(_A					,      vocab_size=2					,      rng=_A      )
    # make sure that at least one token is attended to for each batch
    _UpperCAmelCase          =   1
    return attn_mask




@require_flax
class 		a_						:
    a      :				Union[str, Any]     =				None
    a      :				List[str]     =				()







    def      _snake_case					(			self			:       int     )		->Optional[Any]:
        _UpperCAmelCase    ,_UpperCAmelCase          =   self.model_tester.prepare_config_and_inputs_for_common()

        # cut to half length & take max batch_size 3
        _UpperCAmelCase          =   2
        _UpperCAmelCase          =   inputs["""input_ids"""].shape[-1] // 2
        _UpperCAmelCase          =   inputs["""input_ids"""][:max_batch_size, :sequence_length]

        _UpperCAmelCase          =   jnp.ones_like(__UpperCamelCase     )
        _UpperCAmelCase          =   attention_mask[:max_batch_size, :sequence_length]

        # generate max 5 tokens
        _UpperCAmelCase          =   input_ids.shape[-1] + 5
        if config.eos_token_id is not None and config.pad_token_id is None:
            # hack to allow generate for models such as GPT2 as is done in `generate()`
            _UpperCAmelCase          =   config.eos_token_id
        return config, input_ids, attention_mask, max_length







    @is_pt_flax_cross_test
    def      _snake_case					(			self			:       Any     )		->List[str]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   False
        _UpperCAmelCase          =   max_length
        _UpperCAmelCase          =   0

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model_class.__name__[4:]  # Skip the "Flax" at the beginning
            _UpperCAmelCase          =   getattr(__UpperCamelCase     ,					__UpperCamelCase     )
            _UpperCAmelCase          =   pt_model_class(__UpperCamelCase     ).eval()
            _UpperCAmelCase          =   load_flax_weights_in_pytorch_model(__UpperCamelCase     ,					flax_model.params     )

            _UpperCAmelCase          =   flax_model.generate(__UpperCamelCase     ).sequences
            _UpperCAmelCase          =   pt_model.generate(torch.tensor(__UpperCamelCase     ,					dtype=torch.long     )     )

            if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
                _UpperCAmelCase          =   flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

            self.assertListEqual(pt_generation_outputs.numpy().tolist()     ,					flax_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       List[Any]     )		->List[Any]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   False
        _UpperCAmelCase          =   max_length

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       List[str]     )		->List[str]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   True
        _UpperCAmelCase          =   max_length

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       Any     )		->List[str]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   False
        _UpperCAmelCase          =   max_length
        _UpperCAmelCase          =   2

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       Tuple     )		->Dict:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   False
        _UpperCAmelCase          =   max_length
        _UpperCAmelCase          =   2
        _UpperCAmelCase          =   2

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[0]     ,					input_ids.shape[0] * config.num_return_sequences     )







    def      _snake_case					(			self			:       Dict     )		->List[str]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   True
        _UpperCAmelCase          =   max_length
        _UpperCAmelCase          =   0.8
        _UpperCAmelCase          =   10
        _UpperCAmelCase          =   0.3
        _UpperCAmelCase          =   1
        _UpperCAmelCase          =   8
        _UpperCAmelCase          =   9

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       Dict     )		->List[Any]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   max_length
        _UpperCAmelCase          =   1
        _UpperCAmelCase          =   8
        _UpperCAmelCase          =   9

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       Optional[Any]     )		->Optional[Any]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()
        _UpperCAmelCase          =   max_length
        _UpperCAmelCase          =   2
        _UpperCAmelCase          =   1
        _UpperCAmelCase          =   8
        _UpperCAmelCase          =   9

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       str     )		->Optional[int]:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()

        # pad attention mask on the left
        _UpperCAmelCase          =   attention_mask.at[(0, 0)].set(0     )

        _UpperCAmelCase          =   False
        _UpperCAmelCase          =   max_length

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ,					attention_mask=__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ,					attention_mask=__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       List[str]     )		->Tuple:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()

        # pad attention mask on the left
        _UpperCAmelCase          =   attention_mask.at[(0, 0)].set(0     )

        _UpperCAmelCase          =   True
        _UpperCAmelCase          =   max_length

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ,					attention_mask=__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ,					attention_mask=__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







    def      _snake_case					(			self			:       List[Any]     )		->Dict:
        _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase          =   self._get_input_ids_and_config()

        # pad attention mask on the left
        _UpperCAmelCase          =   attention_mask.at[(0, 0)].set(0     )

        _UpperCAmelCase          =   2
        _UpperCAmelCase          =   max_length

        for model_class in self.all_generative_model_classes:
            _UpperCAmelCase          =   model_class(__UpperCamelCase     )

            _UpperCAmelCase          =   model.generate(__UpperCamelCase     ,					attention_mask=__UpperCamelCase     ).sequences
            self.assertEqual(generation_outputs.shape[-1]     ,					__UpperCamelCase     )

            _UpperCAmelCase          =   jit(model.generate     )
            _UpperCAmelCase          =   jit_generate(__UpperCamelCase     ,					attention_mask=__UpperCamelCase     ).sequences

            self.assertListEqual(generation_outputs.tolist()     ,					jit_generation_outputs.tolist()     )







@require_flax
class 		a_						( unittest.TestCase	):







    def      _snake_case					(			self			:       int     )		->int:
        _UpperCAmelCase          =   AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-bert"""     )
        _UpperCAmelCase          =   FlaxAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""     )

        _UpperCAmelCase          =   """Hello world"""
        _UpperCAmelCase          =   tokenizer(__UpperCamelCase     ,					return_tensors="""np"""     ).input_ids

        # typos are quickly detected (the correct argument is `do_sample`)
        with self.assertRaisesRegex(__UpperCamelCase     ,					"""do_samples"""     ):
            model.generate(__UpperCamelCase     ,					do_samples=__UpperCamelCase     )

        # arbitrary arguments that will not be used anywhere are also not accepted
        with self.assertRaisesRegex(__UpperCamelCase     ,					"""foo"""     ):
            _UpperCAmelCase          =   {"""foo""": """bar"""}
            model.generate(__UpperCamelCase     ,					**__UpperCamelCase     )

"""simple docstring"""



import unittest

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

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()




class 		a_						( unittest.TestCase	):







    @property
    def      _snake_case					(			self			:       Dict     )		->int:






        '''simple docstring'''






        torch.manual_seed(0     )
        _UpperCAmelCase          =   UNetaDModel(
            block_out_channels=(32, 64)     ,					layers_per_block=2     ,					sample_size=32     ,					in_channels=3     ,					out_channels=3     ,					down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")     ,					up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")     ,					)
        return model







    @property
    def      _snake_case					(			self			:       Optional[Any]     )		->str:






        '''simple docstring'''






        torch.manual_seed(0     )
        _UpperCAmelCase          =   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     ,					)
        return model







    @property
    def      _snake_case					(			self			:       List[Any]     )		->Optional[int]:






        '''simple docstring'''






        torch.manual_seed(0     )
        _UpperCAmelCase          =   CLIPTextConfig(
            bos_token_id=0     ,					eos_token_id=2     ,					hidden_size=32     ,					intermediate_size=37     ,					layer_norm_eps=1e-05     ,					num_attention_heads=4     ,					num_hidden_layers=5     ,					pad_token_id=1     ,					vocab_size=10_00     ,					)
        return CLIPTextModel(__UpperCamelCase     )







    def      _snake_case					(			self			:       List[str]     )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase          =   self.dummy_uncond_unet
        _UpperCAmelCase          =   DDIMScheduler()
        _UpperCAmelCase          =   self.dummy_vq_model

        _UpperCAmelCase          =   LDMPipeline(unet=__UpperCamelCase     ,					vqvae=__UpperCamelCase     ,					scheduler=__UpperCamelCase     )
        ldm.to(__UpperCamelCase     )
        ldm.set_progress_bar_config(disable=__UpperCamelCase     )

        _UpperCAmelCase          =   torch.manual_seed(0     )
        _UpperCAmelCase          =   ldm(generator=__UpperCamelCase     ,					num_inference_steps=2     ,					output_type="""numpy"""     ).images

        _UpperCAmelCase          =   torch.manual_seed(0     )
        _UpperCAmelCase          =   ldm(generator=__UpperCamelCase     ,					num_inference_steps=2     ,					output_type="""numpy"""     ,					return_dict=__UpperCamelCase     )[0]

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

        assert image.shape == (1, 64, 64, 3)
        _UpperCAmelCase          =   np.array([0.8_5_1_2, 0.8_1_8, 0.6_4_1_1, 0.6_8_0_8, 0.4_4_6_5, 0.5_6_1_8, 0.4_6, 0.6_2_3_1, 0.5_1_7_2]     )
        _UpperCAmelCase          =   1e-2 if torch_device != """mps""" else 3e-2

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < tolerance
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < tolerance







@slow
@require_torch
class 		a_						( unittest.TestCase	):







    def      _snake_case					(			self			:       List[Any]     )		->Optional[Any]:






        '''simple docstring'''






        _UpperCAmelCase          =   LDMPipeline.from_pretrained("""CompVis/ldm-celebahq-256"""     )
        ldm.to(__UpperCamelCase     )
        ldm.set_progress_bar_config(disable=__UpperCamelCase     )

        _UpperCAmelCase          =   torch.manual_seed(0     )
        _UpperCAmelCase          =   ldm(generator=__UpperCamelCase     ,					num_inference_steps=5     ,					output_type="""numpy"""     ).images

        _UpperCAmelCase          =   image[0, -3:, -3:, -1]

        assert image.shape == (1, 2_56, 2_56, 3)
        _UpperCAmelCase          =   np.array([0.4_3_9_9, 0.4_4_9_7_5, 0.4_6_8_2_5, 0.4_7_4, 0.4_3_5_9, 0.4_5_8_1, 0.4_5_0_9_5, 0.4_3_4_1, 0.4_4_4_7]     )
        _UpperCAmelCase          =   1e-2 if torch_device != """mps""" else 3e-2

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < tolerance