




"""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

		_lowerCamelCase  :			Union[str, Any]        =       '0.12'  # assumed parallelism: 8


if is_torch_available():
		import torch

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase=None	):




	"""simple docstring"""






	if rng is None:
		A_  :				Optional[int]	   =		random.Random()

	A_  :				Optional[int]	   =		1
	for dim in shape:
		total_dims *= dim

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

	A_  :				int	   =		np.array(_UpperCAmelCase  ,      dtype=jnp.intaa	).reshape(_UpperCAmelCase	)

	return output

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase=None	):




	"""simple docstring"""






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

@require_flax
class        lowercase      :
				__lowerCAmelCase    :    Any         =   None
				__lowerCAmelCase    :    Optional[Any]         =   ()







				def 					a_  (							self   :					Optional[Any]       ):



					"""simple docstring"""

					A_   , A_  :				int	   =		self.model_tester.prepare_config_and_inputs_for_common()

					# cut to half length & take max batch_size 3
					A_  :				str	   =		2
					A_  :				str	   =		inputs['''input_ids'''].shape[-1] // 2
					A_  :				Union[str, Any]	   =		inputs['''input_ids'''][:max_batch_size, :sequence_length]

					A_  :				Optional[Any]	   =		jnp.ones_like(_lowerCamelCase       )
					A_  :				Any	   =		attention_mask[:max_batch_size, :sequence_length]

					# generate max 5 tokens
					A_  :				List[str]	   =		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()`
						A_  :				Tuple	   =		config.eos_token_id
					return config, input_ids, attention_mask, max_length







				@is_pt_flax_cross_test
				def 					a_  (							self   :					Optional[Any]       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				List[str]	   =		self._get_input_ids_and_config()
					A_  :				int	   =		False
					A_  :				int	   =		max_length
					A_  :				str	   =		0

					for model_class in self.all_generative_model_classes:
						A_  :				Dict	   =		model_class(_lowerCamelCase       )

						A_  :				Tuple	   =		model_class.__name__[4:]  # Skip the "Flax" at the beginning
						A_  :				str	   =		getattr(_lowerCamelCase      ,     _lowerCamelCase       )
						A_  :				List[str]	   =		pt_model_class(_lowerCamelCase       ).eval()
						A_  :				Optional[int]	   =		load_flax_weights_in_pytorch_model(_lowerCamelCase      ,     flax_model.params       )

						A_  :				str	   =		flax_model.generate(_lowerCamelCase       ).sequences
						A_  :				Dict	   =		pt_model.generate(torch.tensor(_lowerCamelCase      ,     dtype=torch.long       )       )

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

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







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Optional[int]	   =		self._get_input_ids_and_config()
					A_  :				List[Any]	   =		False
					A_  :				Tuple	   =		max_length

					for model_class in self.all_generative_model_classes:
						A_  :				str	   =		model_class(_lowerCamelCase       )

						A_  :				Optional[Any]	   =		model.generate(_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				Tuple	   =		jit(model.generate       )
						A_  :				List[str]	   =		jit_generate(_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Dict	   =		self._get_input_ids_and_config()
					A_  :				Any	   =		True
					A_  :				Dict	   =		max_length

					for model_class in self.all_generative_model_classes:
						A_  :				Tuple	   =		model_class(_lowerCamelCase       )

						A_  :				Optional[int]	   =		model.generate(_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				int	   =		jit(model.generate       )
						A_  :				Union[str, Any]	   =		jit_generate(_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					Optional[Any]       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				List[str]	   =		self._get_input_ids_and_config()
					A_  :				List[str]	   =		False
					A_  :				str	   =		max_length
					A_  :				int	   =		2

					for model_class in self.all_generative_model_classes:
						A_  :				Union[str, Any]	   =		model_class(_lowerCamelCase       )

						A_  :				Optional[Any]	   =		model.generate(_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				List[Any]	   =		jit(model.generate       )
						A_  :				List[Any]	   =		jit_generate(_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					str       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Optional[Any]	   =		self._get_input_ids_and_config()
					A_  :				Union[str, Any]	   =		False
					A_  :				int	   =		max_length
					A_  :				Any	   =		2
					A_  :				Dict	   =		2

					for model_class in self.all_generative_model_classes:
						A_  :				List[str]	   =		model_class(_lowerCamelCase       )

						A_  :				int	   =		model.generate(_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[0]      ,     input_ids.shape[0] * config.num_return_sequences       )







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Optional[Any]	   =		self._get_input_ids_and_config()
					A_  :				List[Any]	   =		True
					A_  :				List[Any]	   =		max_length
					A_  :				Dict	   =		0.8
					A_  :				Union[str, Any]	   =		10
					A_  :				Optional[Any]	   =		0.3
					A_  :				List[str]	   =		1
					A_  :				int	   =		8
					A_  :				str	   =		9

					for model_class in self.all_generative_model_classes:
						A_  :				Union[str, Any]	   =		model_class(_lowerCamelCase       )

						A_  :				Optional[int]	   =		model.generate(_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				int	   =		jit(model.generate       )
						A_  :				Optional[Any]	   =		jit_generate(_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					Optional[Any]       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Optional[Any]	   =		self._get_input_ids_and_config()
					A_  :				str	   =		max_length
					A_  :				str	   =		1
					A_  :				Tuple	   =		8
					A_  :				Optional[Any]	   =		9

					for model_class in self.all_generative_model_classes:
						A_  :				Optional[Any]	   =		model_class(_lowerCamelCase       )

						A_  :				Tuple	   =		model.generate(_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				List[str]	   =		jit(model.generate       )
						A_  :				Union[str, Any]	   =		jit_generate(_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					Optional[Any]       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Union[str, Any]	   =		self._get_input_ids_and_config()
					A_  :				Any	   =		max_length
					A_  :				List[str]	   =		2
					A_  :				List[str]	   =		1
					A_  :				List[str]	   =		8
					A_  :				List[str]	   =		9

					for model_class in self.all_generative_model_classes:
						A_  :				List[str]	   =		model_class(_lowerCamelCase       )

						A_  :				int	   =		model.generate(_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				Tuple	   =		jit(model.generate       )
						A_  :				Tuple	   =		jit_generate(_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					int       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Optional[Any]	   =		self._get_input_ids_and_config()

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

					A_  :				str	   =		False
					A_  :				int	   =		max_length

					for model_class in self.all_generative_model_classes:
						A_  :				int	   =		model_class(_lowerCamelCase       )

						A_  :				str	   =		model.generate(_lowerCamelCase      ,     attention_mask=_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				Any	   =		jit(model.generate       )
						A_  :				Union[str, Any]	   =		jit_generate(_lowerCamelCase      ,     attention_mask=_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					int       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				Optional[Any]	   =		self._get_input_ids_and_config()

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

					A_  :				str	   =		True
					A_  :				Optional[int]	   =		max_length

					for model_class in self.all_generative_model_classes:
						A_  :				Union[str, Any]	   =		model_class(_lowerCamelCase       )

						A_  :				Union[str, Any]	   =		model.generate(_lowerCamelCase      ,     attention_mask=_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				Optional[Any]	   =		jit(model.generate       )
						A_  :				Optional[int]	   =		jit_generate(_lowerCamelCase      ,     attention_mask=_lowerCamelCase       ).sequences

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







				def 					a_  (							self   :					List[str]       ):



					"""simple docstring"""

					A_   , A_   , A_   , A_  :				List[Any]	   =		self._get_input_ids_and_config()

					# pad attention mask on the left
					A_  :				Union[str, Any]	   =		attention_mask.at[(0, 0)].set(0       )

					A_  :				Optional[Any]	   =		2
					A_  :				List[str]	   =		max_length

					for model_class in self.all_generative_model_classes:
						A_  :				str	   =		model_class(_lowerCamelCase       )

						A_  :				Tuple	   =		model.generate(_lowerCamelCase      ,     attention_mask=_lowerCamelCase       ).sequences
						self.assertEqual(generation_outputs.shape[-1]      ,     _lowerCamelCase       )

						A_  :				Optional[int]	   =		jit(model.generate       )
						A_  :				str	   =		jit_generate(_lowerCamelCase      ,     attention_mask=_lowerCamelCase       ).sequences

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

@require_flax
class        lowercase      (    unittest.TestCase):







				def 					a_  (							self   :					Union[str, Any]       ):



					"""simple docstring"""

					A_  :				Any	   =		AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''       )
					A_  :				Any	   =		FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''       )

					A_  :				Optional[Any]	   =		'''Hello world'''
					A_  :				Any	   =		tokenizer(_lowerCamelCase      ,     return_tensors='''np'''       ).input_ids

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

					# arbitrary arguments that will not be used anywhere are also not accepted
					with self.assertRaisesRegex(_lowerCamelCase      ,     '''foo'''       ):
						A_  :				Optional[int]	   =		{'''foo''': '''bar'''}
						model.generate(_lowerCamelCase      ,     **_lowerCamelCase       )











"""simple docstring"""

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


_lowerCamelCase  :			Any        =       False

class        lowercase      (    unittest.TestCase):







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

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







				@property
				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

					return 12







				@property
				def 					a_  (							self   :					List[str]       ):



					"""simple docstring"""

					return 12







				@property
				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					return 32







				@property
				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

					torch.manual_seed(0       )
					A_  :				Union[str, 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 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					A_  :				int	   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )
					return tokenizer







				@property
				def 					a_  (							self   :					int       ):



					"""simple docstring"""

					torch.manual_seed(0       )
					A_  :				Optional[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=10_00      ,     )
					return CLIPTextModel(_lowerCamelCase       )







				@property
				def 					a_  (							self   :					Optional[int]       ):



					"""simple docstring"""

					torch.manual_seed(0       )

					A_  :				Optional[Any]	   =		12
					A_  :				Optional[int]	   =		12

					A_  :				int	   =		{
					    '''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''',
					}

					A_  :				Tuple	   =		TransformeraDModel(**_lowerCamelCase       )
					return model







				def 					a_  (							self   :					Optional[int]       ):



					"""simple docstring"""

					A_  :				Union[str, Any]	   =		'''cpu'''

					A_  :				Union[str, Any]	   =		self.dummy_vqvae
					A_  :				str	   =		self.dummy_text_encoder
					A_  :				List[Any]	   =		self.dummy_tokenizer
					A_  :				int	   =		self.dummy_transformer
					A_  :				Any	   =		VQDiffusionScheduler(self.num_embed       )
					A_  :				Optional[Any]	   =		LearnedClassifierFreeSamplingEmbeddings(learnable=_lowerCamelCase       )

					A_  :				Dict	   =		VQDiffusionPipeline(
					    vqvae=_lowerCamelCase      ,     text_encoder=_lowerCamelCase      ,     tokenizer=_lowerCamelCase      ,     transformer=_lowerCamelCase      ,     scheduler=_lowerCamelCase      ,     learned_classifier_free_sampling_embeddings=_lowerCamelCase      ,     )
					A_  :				List[Any]	   =		pipe.to(_lowerCamelCase       )
					pipe.set_progress_bar_config(disable=_lowerCamelCase       )

					A_  :				List[Any]	   =		'''teddy bear playing in the pool'''

					A_  :				List[Any]	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				List[Any]	   =		pipe([prompt]      ,     generator=_lowerCamelCase      ,     num_inference_steps=2      ,     output_type='''np'''       )
					A_  :				Any	   =		output.images

					A_  :				List[str]	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				List[Any]	   =		pipe(
					    [prompt]      ,     generator=_lowerCamelCase      ,     output_type='''np'''      ,     return_dict=_lowerCamelCase      ,     num_inference_steps=2       )[0]

					A_  :				Optional[int]	   =		image[0, -3:, -3:, -1]
					A_  :				Optional[int]	   =		image_from_tuple[0, -3:, -3:, -1]

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

					A_  :				Optional[int]	   =		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 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					A_  :				Union[str, Any]	   =		'''cpu'''

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

					A_  :				int	   =		VQDiffusionPipeline(
					    vqvae=_lowerCamelCase      ,     text_encoder=_lowerCamelCase      ,     tokenizer=_lowerCamelCase      ,     transformer=_lowerCamelCase      ,     scheduler=_lowerCamelCase      ,     learned_classifier_free_sampling_embeddings=_lowerCamelCase      ,     )
					A_  :				List[Any]	   =		pipe.to(_lowerCamelCase       )
					pipe.set_progress_bar_config(disable=_lowerCamelCase       )

					A_  :				Any	   =		'''teddy bear playing in the pool'''

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

					A_  :				Optional[int]	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				List[str]	   =		pipe(
					    [prompt]      ,     generator=_lowerCamelCase      ,     output_type='''np'''      ,     return_dict=_lowerCamelCase      ,     num_inference_steps=2       )[0]

					A_  :				Optional[int]	   =		image[0, -3:, -3:, -1]
					A_  :				Optional[Any]	   =		image_from_tuple[0, -3:, -3:, -1]

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

					A_  :				str	   =		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        lowercase      (    unittest.TestCase):







				def 					a_  (							self   :					Any       ):



					"""simple docstring"""

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







				def 					a_  (							self   :					str       ):



					"""simple docstring"""

					A_  :				str	   =		load_numpy(
					    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
					    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''       )

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

					# requires GPU generator for gumbel softmax
					# don't use GPU generator in tests though
					A_  :				Dict	   =		torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )
					A_  :				Union[str, Any]	   =		pipeline(
					    '''teddy bear playing in the pool'''      ,     num_images_per_prompt=1      ,     generator=_lowerCamelCase      ,     output_type='''np'''      ,     )

					A_  :				Optional[int]	   =		output.images[0]

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





