






from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers.models.xglm.modeling_tf_xglm import (
		    TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFXGLMForCausalLM,
		    TFXGLMModel,
		)





@require_tf
class 	snake_case__ :
	lowercase__						:					Dict											=   XGLMConfig
	lowercase__						:					str											=   {}
	lowercase__						:					List[str]											=   '''gelu'''
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__=14      ,       lowerCAmelCase__=7      ,       lowerCAmelCase__=True      ,       lowerCAmelCase__=True      ,       lowerCAmelCase__=True      ,       lowerCAmelCase__=99      ,       lowerCAmelCase__=32      ,       lowerCAmelCase__=2      ,       lowerCAmelCase__=4      ,       lowerCAmelCase__=37      ,       lowerCAmelCase__="gelu"      ,       lowerCAmelCase__=0.1      ,       lowerCAmelCase__=0.1      ,       lowerCAmelCase__=5_12      ,       lowerCAmelCase__=0.0_2      ,       )							->				Optional[int]:
							__magic_name__      :       List[str]       						= parent
							__magic_name__      :       Dict       						= batch_size
							__magic_name__      :       Any       						= seq_length
							__magic_name__      :       Dict       						= is_training
							__magic_name__      :       Dict       						= use_input_mask
							__magic_name__      :       Dict       						= use_labels
							__magic_name__      :       int       						= vocab_size
							__magic_name__      :       List[Any]       						= d_model
							__magic_name__      :       List[str]       						= num_hidden_layers
							__magic_name__      :       int       						= num_attention_heads
							__magic_name__      :       Optional[int]       						= ffn_dim
							__magic_name__      :       Dict       						= activation_function
							__magic_name__      :       str       						= activation_dropout
							__magic_name__      :       Tuple       						= attention_dropout
							__magic_name__      :       Union[str, Any]       						= max_position_embeddings
							__magic_name__      :       Union[str, Any]       						= initializer_range
							__magic_name__      :       Any       						= None
							__magic_name__      :       Tuple       						= 0
							__magic_name__      :       str       						= 2
							__magic_name__      :       Any       						= 1
	def       __magic_name__    (    self			)							->				Any:
							return XGLMConfig.from_pretrained("""facebook/xglm-564M"""			)
	def       __magic_name__    (    self			)							->				int:
							__magic_name__      :       List[Any]       						= tf.clip_by_value(
							    ids_tensor([self.batch_size, self.seq_length]      ,       self.vocab_size			)      ,       clip_value_min=0      ,       clip_value_max=3			)

							__magic_name__      :       Union[str, Any]       						= None
							if self.use_input_mask:
													__magic_name__      :       Optional[int]       						= random_attention_mask([self.batch_size, self.seq_length]			)

							__magic_name__      :       List[str]       						= self.get_config()

							__magic_name__      :       Dict       						= floats_tensor([self.num_hidden_layers, self.num_attention_heads]      ,       2			)

							return (
							    config,
							    input_ids,
							    input_mask,
							    head_mask,
							)
	def       __magic_name__    (    self			)							->				Optional[Any]:
							return XGLMConfig(
							    vocab_size=self.vocab_size      ,       d_model=self.hidden_size      ,       num_layers=self.num_hidden_layers      ,       attention_heads=self.num_attention_heads      ,       ffn_dim=self.ffn_dim      ,       activation_function=self.activation_function      ,       activation_dropout=self.activation_dropout      ,       attention_dropout=self.attention_dropout      ,       max_position_embeddings=self.max_position_embeddings      ,       initializer_range=self.initializer_range      ,       use_cache=lowerCAmelCase__      ,       bos_token_id=self.bos_token_id      ,       eos_token_id=self.eos_token_id      ,       pad_token_id=self.pad_token_id      ,       return_dict=lowerCAmelCase__      ,       )







	def       __magic_name__    (    self			)							->				Optional[Any]:
							__magic_name__      :       int       						= self.prepare_config_and_inputs()

							(
							    (
							    __magic_name__
							)						,(
							    __magic_name__
							)						,(
							    __magic_name__
							)						,(
							    __magic_name__
							)						,
							)      :       List[Any]       						= config_and_inputs

							__magic_name__      :       Tuple       						= {
							    """input_ids""": input_ids,
							    """head_mask""": head_mask,
							}

							return config, inputs_dict





@require_tf
class 	snake_case__ (							_lowerCAmelCase						,  _lowerCAmelCase						,  unittest.TestCase				):
	lowercase__						:					int											=   (TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
	lowercase__						:					Any											=   (TFXGLMForCausalLM,) if is_tf_available() else ()
	lowercase__						:					Optional[Any]											=   (
	    {'''feature-extraction''': TFXGLMModel, '''text-generation''': TFXGLMForCausalLM} if is_tf_available() else {}
	)
	lowercase__						:					List[Any]											=   False
	lowercase__						:					Any											=   False
	lowercase__						:					Optional[int]											=   False
	def       __magic_name__    (    self			)							->				List[str]:
							__magic_name__      :       List[str]       						= TFXGLMModelTester(self			)
							__magic_name__      :       List[Any]       						= ConfigTester(self      ,       config_class=lowerCAmelCase__      ,       n_embd=37			)
	def       __magic_name__    (    self			)							->				str:
							self.config_tester.run_common_tests()
	@slow
	def       __magic_name__    (    self			)							->				List[str]:
							for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__magic_name__      :       Dict       						= TFXGLMModel.from_pretrained(lowerCAmelCase__			)
													self.assertIsNotNone(lowerCAmelCase__			)







	@unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""			)
	def       __magic_name__    (    self			)							->				int:
							super().test_resize_token_embeddings()





@require_tf
class 	snake_case__ (							unittest.TestCase				):
	@slow
	def       __magic_name__    (    self      ,       lowerCAmelCase__=True			)							->				int:
							__magic_name__      :       str       						= TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""			)
							__magic_name__      :       int       						= tf.convert_to_tensor([[2, 2_68, 98_65]]      ,       dtype=tf.intaa			)  # The dog
							# </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
							# fmt: off
							__magic_name__      :       Optional[int]       						= [2, 2_68, 98_65, 67, 11, 19_88, 5_72_52, 98_65, 5, 9_84, 67, 19_88, 21_38_38, 16_58, 53, 7_04_46, 33, 66_57, 2_78, 15_81]
							# fmt: on
							__magic_name__      :       Optional[int]       						= model.generate(lowerCAmelCase__      ,       do_sample=lowerCAmelCase__      ,       num_beams=1			)
							if verify_outputs:
													self.assertListEqual(output_ids[0].numpy().tolist()      ,       lowerCAmelCase__			)
	@slow
	def       __magic_name__    (    self			)							->				Optional[int]:
							__magic_name__      :       Dict       						= XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""			)
							__magic_name__      :       Dict       						= TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""			)

							tf.random.set_seed(0			)
							__magic_name__      :       List[Any]       						= tokenizer("""Today is a nice day and"""      ,       return_tensors="""tf"""			)
							__magic_name__      :       Optional[int]       						= tokenized.input_ids
							# forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
							with tf.device(""":/CPU:0"""			):
													__magic_name__      :       int       						= model.generate(lowerCAmelCase__      ,       do_sample=lowerCAmelCase__      ,       seed=[7, 0]			)
							__magic_name__      :       Tuple       						= tokenizer.decode(output_ids[0]      ,       skip_special_tokens=lowerCAmelCase__			)

							__magic_name__      :       Union[str, Any]       						= (
							    """Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due"""
							)
							self.assertEqual(lowerCAmelCase__      ,       lowerCAmelCase__			)







	@slow
	def       __magic_name__    (    self			)							->				Any:
							__magic_name__      :       Dict       						= TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""			)
							__magic_name__      :       str       						= XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""			)

							__magic_name__      :       Tuple       						= """left"""

							# use different length sentences to test batching
							__magic_name__      :       Optional[int]       						= [
							    """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
							    """left-padding, such as in batched generation. The output for the sequence below should be the same """
							    """regardless of whether left padding is applied or not. When""",
							    """Hello, my dog is a little""",
							]

							__magic_name__      :       int       						= tokenizer(lowerCAmelCase__      ,       return_tensors="""tf"""      ,       padding=lowerCAmelCase__			)
							__magic_name__      :       Dict       						= inputs["""input_ids"""]

							__magic_name__      :       List[str]       						= model.generate(input_ids=lowerCAmelCase__      ,       attention_mask=inputs["""attention_mask"""]      ,       max_new_tokens=12			)

							__magic_name__      :       Any       						= tokenizer(sentences[0]      ,       return_tensors="""tf"""			).input_ids
							__magic_name__      :       List[Any]       						= model.generate(input_ids=lowerCAmelCase__      ,       max_new_tokens=12			)

							__magic_name__      :       Dict       						= tokenizer(sentences[1]      ,       return_tensors="""tf"""			).input_ids
							__magic_name__      :       str       						= model.generate(input_ids=lowerCAmelCase__      ,       max_new_tokens=12			)

							__magic_name__      :       Any       						= tokenizer.batch_decode(lowerCAmelCase__      ,       skip_special_tokens=lowerCAmelCase__			)
							__magic_name__      :       str       						= tokenizer.decode(output_non_padded[0]      ,       skip_special_tokens=lowerCAmelCase__			)
							__magic_name__      :       str       						= tokenizer.decode(output_padded[0]      ,       skip_special_tokens=lowerCAmelCase__			)

							__magic_name__      :       Any       						= [
							    """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
							    """left-padding, such as in batched generation. The output for the sequence below should be the same """
							    """regardless of whether left padding is applied or not. When left padding is applied, the sequence will be """
							    """a single""",
							    """Hello, my dog is a little bit of a shy one, but he is very friendly""",
							]
							self.assertListEqual(lowerCAmelCase__      ,       lowerCAmelCase__			)
							self.assertListEqual(lowerCAmelCase__      ,       [non_padded_sentence, padded_sentence]			)








import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class 	snake_case__ (							_lowerCAmelCase						,  _lowerCAmelCase						,  _lowerCAmelCase						,  unittest.TestCase				):
	lowercase__						:					str											=   StableUnCLIPImgaImgPipeline
	lowercase__						:					Any											=   TEXT_GUIDED_IMAGE_VARIATION_PARAMS
	lowercase__						:					Tuple											=   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	lowercase__						:					Union[str, Any]											=   frozenset(
	    []				)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
	lowercase__						:					Union[str, Any]											=   frozenset([]				)
	def       __magic_name__    (    self			)							->				Tuple:
							__magic_name__      :       Any       						= 32
							__magic_name__      :       Union[str, Any]       						= embedder_hidden_size

							# image encoding components

							__magic_name__      :       Optional[int]       						= CLIPImageProcessor(crop_size=32      ,       size=32			)

							torch.manual_seed(0			)
							__magic_name__      :       Optional[Any]       						= CLIPVisionModelWithProjection(
							    CLIPVisionConfig(
							        hidden_size=lowerCAmelCase__      ,       projection_dim=lowerCAmelCase__      ,       num_hidden_layers=5      ,       num_attention_heads=4      ,       image_size=32      ,       intermediate_size=37      ,       patch_size=1      ,       )			)

							# regular denoising components

							torch.manual_seed(0			)
							__magic_name__      :       Any       						= StableUnCLIPImageNormalizer(embedding_dim=lowerCAmelCase__			)
							__magic_name__      :       int       						= DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""			)

							torch.manual_seed(0			)
							__magic_name__      :       Tuple       						= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

							torch.manual_seed(0			)
							__magic_name__      :       Optional[Any]       						= CLIPTextModel(
							    CLIPTextConfig(
							        bos_token_id=0      ,       eos_token_id=2      ,       hidden_size=lowerCAmelCase__      ,       projection_dim=32      ,       intermediate_size=37      ,       layer_norm_eps=1e-0_5      ,       num_attention_heads=4      ,       num_hidden_layers=5      ,       pad_token_id=1      ,       vocab_size=10_00      ,       )			)

							torch.manual_seed(0			)
							__magic_name__      :       str       						= UNetaDConditionModel(
							    sample_size=32      ,       in_channels=4      ,       out_channels=4      ,       down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")      ,       up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")      ,       block_out_channels=(32, 64)      ,       attention_head_dim=(2, 4)      ,       class_embed_type="""projection"""      ,       projection_class_embeddings_input_dim=embedder_projection_dim * 2      ,       cross_attention_dim=lowerCAmelCase__      ,       layers_per_block=1      ,       upcast_attention=lowerCAmelCase__      ,       use_linear_projection=lowerCAmelCase__      ,       )

							torch.manual_seed(0			)
							__magic_name__      :       Optional[Any]       						= DDIMScheduler(
							    beta_schedule="""scaled_linear"""      ,       beta_start=0.0_0_0_8_5      ,       beta_end=0.0_1_2      ,       prediction_type="""v_prediction"""      ,       set_alpha_to_one=lowerCAmelCase__      ,       steps_offset=1      ,       )

							torch.manual_seed(0			)
							__magic_name__      :       List[str]       						= AutoencoderKL()

							__magic_name__      :       List[str]       						= {
							    # image encoding components
							    """feature_extractor""": feature_extractor,
							    """image_encoder""": image_encoder.eval(),
							    # image noising components
							    """image_normalizer""": image_normalizer.eval(),
							    """image_noising_scheduler""": image_noising_scheduler,
							    # regular denoising components
							    """tokenizer""": tokenizer,
							    """text_encoder""": text_encoder.eval(),
							    """unet""": unet.eval(),
							    """scheduler""": scheduler,
							    """vae""": vae.eval(),
							}

							return components
	def       __magic_name__    (    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__=0      ,       lowerCAmelCase__=True			)							->				List[Any]:
							if str(lowerCAmelCase__			).startswith("""mps"""			):
													__magic_name__      :       Optional[int]       						= torch.manual_seed(lowerCAmelCase__			)
							else:
													__magic_name__      :       Optional[int]       						= torch.Generator(device=lowerCAmelCase__			).manual_seed(lowerCAmelCase__			)

							__magic_name__      :       Any       						= floats_tensor((1, 3, 32, 32)      ,       rng=random.Random(lowerCAmelCase__			)			).to(lowerCAmelCase__			)

							if pil_image:
													__magic_name__      :       Optional[Any]       						= input_image * 0.5 + 0.5
													__magic_name__      :       int       						= input_image.clamp(0      ,       1			)
													__magic_name__      :       Union[str, Any]       						= input_image.cpu().permute(0      ,       2      ,       3      ,       1			).float().numpy()
													__magic_name__      :       Optional[int]       						= DiffusionPipeline.numpy_to_pil(lowerCAmelCase__			)[0]

							return {
							    "prompt": "An anime racoon running a marathon",
							    "image": input_image,
							    "generator": generator,
							    "num_inference_steps": 2,
							    "output_type": "np",
							}
	@skip_mps
	def       __magic_name__    (    self			)							->				List[Any]:
							__magic_name__      :       Optional[Any]       						= """cpu"""  # ensure determinism for the device-dependent torch.Generator
							__magic_name__      :       List[str]       						= self.get_dummy_components()
							__magic_name__      :       int       						= StableUnCLIPImgaImgPipeline(**lowerCAmelCase__			)
							__magic_name__      :       List[str]       						= sd_pipe.to(lowerCAmelCase__			)
							sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__			)

							__magic_name__      :       Any       						= self.get_dummy_inputs(lowerCAmelCase__			)
							inputs.update({"""image_embeds""": None}			)
							__magic_name__      :       List[str]       						= sd_pipe(**lowerCAmelCase__			).images
							__magic_name__      :       Dict       						= image[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							__magic_name__      :       int       						= np.array([0.3_8_7_2, 0.7_2_2_4, 0.5_6_0_1, 0.4_7_4_1, 0.6_8_7_2, 0.5_8_1_4, 0.4_6_3_6, 0.3_8_6_7, 0.5_0_7_8]			)

							assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-3
	def       __magic_name__    (    self			)							->				Dict:
							__magic_name__      :       int       						= torch_device in ["""cpu""", """mps"""]

							self._test_attention_slicing_forward_pass(test_max_difference=lowerCAmelCase__			)
	def       __magic_name__    (    self			)							->				Union[str, Any]:
							__magic_name__      :       Tuple       						= torch_device in ["""cpu""", """mps"""]

							self._test_inference_batch_single_identical(test_max_difference=lowerCAmelCase__			)







	@unittest.skipIf(
	    torch_device != """cuda""" or not is_xformers_available()      ,       reason="""XFormers attention is only available with CUDA and `xformers` installed"""      ,       )
	def       __magic_name__    (    self			)							->				str:
							self._test_xformers_attention_forwardGenerator_pass(test_max_difference=lowerCAmelCase__			)





@slow
@require_torch_gpu
class 	snake_case__ (							unittest.TestCase				):
	def       __magic_name__    (    self			)							->				int:
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()
	def       __magic_name__    (    self			)							->				List[Any]:
							__magic_name__      :       Optional[Any]       						= load_image(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""			)

							__magic_name__      :       int       						= load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy"""			)

							__magic_name__      :       List[Any]       						= StableUnCLIPImgaImgPipeline.from_pretrained(
							    """fusing/stable-unclip-2-1-l-img2img"""      ,       torch_dtype=torch.floataa			)
							pipe.to(lowerCAmelCase__			)
							pipe.set_progress_bar_config(disable=lowerCAmelCase__			)
							# stable unclip will oom when integration tests are run on a V100,
							# so turn on memory savings
							pipe.enable_attention_slicing()
							pipe.enable_sequential_cpu_offload()

							__magic_name__      :       Union[str, Any]       						= torch.Generator(device="""cpu"""			).manual_seed(0			)
							__magic_name__      :       Union[str, Any]       						= pipe(lowerCAmelCase__      ,       """anime turle"""      ,       generator=lowerCAmelCase__      ,       output_type="""np"""			)

							__magic_name__      :       Optional[Any]       						= output.images[0]

							assert image.shape == (7_68, 7_68, 3)

							assert_mean_pixel_difference(lowerCAmelCase__      ,       lowerCAmelCase__			)
	def       __magic_name__    (    self			)							->				List[Any]:
							__magic_name__      :       Union[str, Any]       						= load_image(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""			)

							__magic_name__      :       Dict       						= load_numpy(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy"""			)

							__magic_name__      :       Union[str, Any]       						= StableUnCLIPImgaImgPipeline.from_pretrained(
							    """fusing/stable-unclip-2-1-h-img2img"""      ,       torch_dtype=torch.floataa			)
							pipe.to(lowerCAmelCase__			)
							pipe.set_progress_bar_config(disable=lowerCAmelCase__			)
							# stable unclip will oom when integration tests are run on a V100,
							# so turn on memory savings
							pipe.enable_attention_slicing()
							pipe.enable_sequential_cpu_offload()

							__magic_name__      :       List[str]       						= torch.Generator(device="""cpu"""			).manual_seed(0			)
							__magic_name__      :       Any       						= pipe(lowerCAmelCase__      ,       """anime turle"""      ,       generator=lowerCAmelCase__      ,       output_type="""np"""			)

							__magic_name__      :       Any       						= output.images[0]

							assert image.shape == (7_68, 7_68, 3)

							assert_mean_pixel_difference(lowerCAmelCase__      ,       lowerCAmelCase__			)







	def       __magic_name__    (    self			)							->				str:
							__magic_name__      :       Any       						= load_image(
							    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"""			)

							torch.cuda.empty_cache()
							torch.cuda.reset_max_memory_allocated()
							torch.cuda.reset_peak_memory_stats()

							__magic_name__      :       Union[str, Any]       						= StableUnCLIPImgaImgPipeline.from_pretrained(
							    """fusing/stable-unclip-2-1-h-img2img"""      ,       torch_dtype=torch.floataa			)
							__magic_name__      :       int       						= pipe.to(lowerCAmelCase__			)
							pipe.set_progress_bar_config(disable=lowerCAmelCase__			)
							pipe.enable_attention_slicing()
							pipe.enable_sequential_cpu_offload()

							__magic_name__      :       List[Any]       						= pipe(
							    lowerCAmelCase__      ,       """anime turtle"""      ,       num_inference_steps=2      ,       output_type="""np"""      ,       )

							__magic_name__      :       List[str]       						= torch.cuda.max_memory_allocated()
							# make sure that less than 7 GB is allocated
							assert mem_bytes < 7 * 10**9
