

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__						 =			False

class        a__    (       unittest.TestCase      ):
			"""simple docstring"""







			def   UpperCamelCase		(     self							)	->    Optional[int]:
							'''simple docstring'''






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


			@property
			def   UpperCamelCase		(     self							)	->    Optional[int]:
							'''simple docstring'''






							return 12


			@property
			def   UpperCamelCase		(     self							)	->    List[Any]:
							'''simple docstring'''






							return 12


			@property
			def   UpperCamelCase		(     self							)	->    Optional[Any]:
							'''simple docstring'''






							return 32


			@property
			def   UpperCamelCase		(     self							)	->    List[str]:
							'''simple docstring'''






							torch.manual_seed(0							)
							A__          =      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							)	->    Dict:
							'''simple docstring'''






							A__          =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)
							return tokenizer


			@property
			def   UpperCamelCase		(     self							)	->    List[str]:
							'''simple docstring'''






							torch.manual_seed(0							)
							A__          =      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(__snake_case							)


			@property
			def   UpperCamelCase		(     self							)	->    int:
							'''simple docstring'''






							torch.manual_seed(0							)

							A__          =      12
							A__          =      12

							A__          =      {
							    """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__          =      TransformeraDModel(**__snake_case							)
							return model


			def   UpperCamelCase		(     self							)	->    Optional[Any]:
							'''simple docstring'''






							A__          =      """cpu"""

							A__          =      self.dummy_vqvae
							A__          =      self.dummy_text_encoder
							A__          =      self.dummy_tokenizer
							A__          =      self.dummy_transformer
							A__          =      VQDiffusionScheduler(self.num_embed							)
							A__          =      LearnedClassifierFreeSamplingEmbeddings(learnable=__snake_case							)

							A__          =      VQDiffusionPipeline(
							    vqvae=__snake_case   ,   text_encoder=__snake_case   ,   tokenizer=__snake_case   ,   transformer=__snake_case   ,   scheduler=__snake_case   ,   learned_classifier_free_sampling_embeddings=__snake_case   ,   )
							A__          =      pipe.to(__snake_case							)
							pipe.set_progress_bar_config(disable=__snake_case							)

							A__          =      """teddy bear playing in the pool"""

							A__          =      torch.Generator(device=__snake_case							).manual_seed(0							)
							A__          =      pipe([prompt]   ,   generator=__snake_case   ,   num_inference_steps=2   ,   output_type="np"							)
							A__          =      output.images

							A__          =      torch.Generator(device=__snake_case							).manual_seed(0							)
							A__          =      pipe(
							    [prompt]   ,   generator=__snake_case   ,   output_type="np"   ,   return_dict=__snake_case   ,   num_inference_steps=2							)[0]

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

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

							A__          =      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							)	->    List[Any]:
							'''simple docstring'''






							A__          =      """cpu"""

							A__          =      self.dummy_vqvae
							A__          =      self.dummy_text_encoder
							A__          =      self.dummy_tokenizer
							A__          =      self.dummy_transformer
							A__          =      VQDiffusionScheduler(self.num_embed							)
							A__          =      LearnedClassifierFreeSamplingEmbeddings(
							    learnable=__snake_case   ,   hidden_size=self.text_embedder_hidden_size   ,   length=tokenizer.model_max_length							)

							A__          =      VQDiffusionPipeline(
							    vqvae=__snake_case   ,   text_encoder=__snake_case   ,   tokenizer=__snake_case   ,   transformer=__snake_case   ,   scheduler=__snake_case   ,   learned_classifier_free_sampling_embeddings=__snake_case   ,   )
							A__          =      pipe.to(__snake_case							)
							pipe.set_progress_bar_config(disable=__snake_case							)

							A__          =      """teddy bear playing in the pool"""

							A__          =      torch.Generator(device=__snake_case							).manual_seed(0							)
							A__          =      pipe([prompt]   ,   generator=__snake_case   ,   num_inference_steps=2   ,   output_type="np"							)
							A__          =      output.images

							A__          =      torch.Generator(device=__snake_case							).manual_seed(0							)
							A__          =      pipe(
							    [prompt]   ,   generator=__snake_case   ,   output_type="np"   ,   return_dict=__snake_case   ,   num_inference_steps=2							)[0]

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

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

							A__          =      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        a__    (       unittest.TestCase      ):
			"""simple docstring"""







			def   UpperCamelCase		(     self							)	->    str:
							'''simple docstring'''






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


			def   UpperCamelCase		(     self							)	->    Dict:
							'''simple docstring'''






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

							A__          =      VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq"							)
							A__          =      pipeline.to(__snake_case							)
							pipeline.set_progress_bar_config(disable=__snake_case							)

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

							A__          =      output.images[0]

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







import math

import tensorflow as tf
from packaging import version




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						str										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						List[str]										=0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0 )						,  x.dtype ) ))

				return x * cdf




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						str										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						Optional[Any]										=tf.cast(math.pi						,  x.dtype )
				__magic_name__						:						int										=tf.cast(0.0_4_4_7_1_5						,  x.dtype )
				__magic_name__						:						Tuple										=0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi ) * (x + coeff * tf.pow(lowerCamelCase						,  3 )) ))

				return x * cdf




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						Any										=tf.convert_to_tensor(lowerCamelCase )

				return x * tf.tanh(tf.math.softplus(lowerCamelCase ) )




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						Optional[Any]										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						Union[str, Any]										=tf.cast(0.0_4_4_7_1_5						,  x.dtype )
				__magic_name__						:						Tuple										=tf.cast(0.7_9_7_8_8_4_5_6_0_8						,  x.dtype )

				return 0.5 * x * (1.0 + tf.tanh(x * coeffa * (1.0 + coeffa * x * x) ))




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						List[str]										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						Dict										=tf.cast(1.7_0_2						,  x.dtype )
				return x * tf.math.sigmoid(coeff * x )




def 		lowerCAmelCase_ (						lowerCamelCase ):
				return tf.clip_by_value(_gelu(lowerCamelCase )						,  -10						,  10 )




def 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase=-1 ):
				__magic_name__					,			__magic_name__						:						List[Any]										=tf.split(lowerCamelCase						,  2						,  axis=lowerCamelCase )
				return a * tf.math.sigmoid(lowerCamelCase )


if version.parse(tf.version.VERSION) >= version.parse("2.4"):




							def 		lowerCAmelCase_ (						lowerCamelCase ):
											return tf.keras.activations.gelu(lowerCamelCase						,  approximate=lowerCamelCase )

							UpperCAmelCase_		:					List[str]          =  tf.keras.activations.gelu



							UpperCAmelCase_		:					Dict          =  approximate_gelu_wrap
else:
							UpperCAmelCase_		:					Dict          =  _gelu
							UpperCAmelCase_		:					str          =  _gelu_new


UpperCAmelCase_		:					Any          =  {
    "gelu": gelu,
    "gelu_10": gelu_aa,
    "gelu_fast": gelu_fast,
    "gelu_new": gelu_new,
    "glu": glu,
    "mish": mish,
    "quick_gelu": quick_gelu,
    "relu": tf.keras.activations.relu,
    "sigmoid": tf.keras.activations.sigmoid,
    "silu": tf.keras.activations.swish,
    "swish": tf.keras.activations.swish,
    "tanh": tf.keras.activations.tanh,
}




def 		lowerCAmelCase_ (						lowerCamelCase ):
				if activation_string in ACTaFN:
								return ACTaFN[activation_string]
				else:
								raise KeyError(F"function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys() )}" )




