
import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     int				)			->						List[Any]:
       SCREAMING_SNAKE_CASE_							:				List[str]		=[tensor.shape for tensor in tensor_list]
       return all(shape == shapes[0] for shape in shapes[1:]				)


class 			lowercase_		(      __lowercase  ,      __lowercase  ,      __lowercase  ,      unittest.TestCase       ):
       __lowerCamelCase				       =						StableDiffusionLatentUpscalePipeline
       __lowerCamelCase				       =						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
           "height",
           "width",
           "cross_attention_kwargs",
           "negative_prompt_embeds",
           "prompt_embeds",
       }
       __lowerCamelCase				       =						PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
       __lowerCamelCase				       =						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
       __lowerCamelCase				       =						frozenset(
           []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
       __lowerCamelCase				       =						frozenset([]       )

       __lowerCamelCase				       =						True


       @property
       def 	_snake_case				(    self				)						->    Tuple:
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=1
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=4
              SCREAMING_SNAKE_CASE_							:				int		=(16, 16)

              SCREAMING_SNAKE_CASE_							:				Any		=floats_tensor((batch_size, num_channels) + sizes						,       rng=random.Random(0				)				).to(_A				)
              return image


       def 	_snake_case				(    self				)						->    Optional[int]:
              torch.manual_seed(0				)
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=UNetaDConditionModel(
                  act_fn='''gelu'''						,       attention_head_dim=8						,       norm_num_groups=_A						,       block_out_channels=[32, 32, 64, 64]						,       time_cond_proj_dim=160						,       conv_in_kernel=1						,       conv_out_kernel=1						,       cross_attention_dim=32						,       down_block_types=(
                      '''KDownBlock2D''',
                      '''KCrossAttnDownBlock2D''',
                      '''KCrossAttnDownBlock2D''',
                      '''KCrossAttnDownBlock2D''',
                  )						,       in_channels=8						,       mid_block_type=_A						,       only_cross_attention=_A						,       out_channels=5						,       resnet_time_scale_shift='''scale_shift'''						,       time_embedding_type='''fourier'''						,       timestep_post_act='''gelu'''						,       up_block_types=('''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KUpBlock2D''')						,       )
              SCREAMING_SNAKE_CASE_							:				List[str]		=AutoencoderKL(
                  block_out_channels=[32, 32, 64, 64]						,       in_channels=3						,       out_channels=3						,       down_block_types=[
                      '''DownEncoderBlock2D''',
                      '''DownEncoderBlock2D''',
                      '''DownEncoderBlock2D''',
                      '''DownEncoderBlock2D''',
                  ]						,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''']						,       latent_channels=4						,       )
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=EulerDiscreteScheduler(prediction_type='''sample'''				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=CLIPTextConfig(
                  bos_token_id=0						,       eos_token_id=2						,       hidden_size=32						,       intermediate_size=37						,       layer_norm_eps=1e-0_5						,       num_attention_heads=4						,       num_hidden_layers=5						,       pad_token_id=1						,       vocab_size=1_000						,       hidden_act='''quick_gelu'''						,       projection_dim=512						,       )
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=CLIPTextModel(_A				)
              SCREAMING_SNAKE_CASE_							:				str		=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

              SCREAMING_SNAKE_CASE_							:				List[Any]		={
                  '''unet''': model.eval(),
                  '''vae''': vae.eval(),
                  '''scheduler''': scheduler,
                  '''text_encoder''': text_encoder,
                  '''tokenizer''': tokenizer,
              }

              return components


       def 	_snake_case				(    self						,       __A						,       __A=0				)						->    Dict:
              if str(_A				).startswith('''mps'''				):
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=torch.manual_seed(_A				)
              else:
                     SCREAMING_SNAKE_CASE_							:				List[Any]		=torch.Generator(device=_A				).manual_seed(_A				)
              SCREAMING_SNAKE_CASE_							:				Any		={
                  '''prompt''': '''A painting of a squirrel eating a burger''',
                  '''image''': self.dummy_image.cpu(),
                  '''generator''': generator,
                  '''num_inference_steps''': 2,
                  '''output_type''': '''numpy''',
              }
              return inputs


       def 	_snake_case				(    self				)						->    Union[str, Any]:
              SCREAMING_SNAKE_CASE_							:				Tuple		='''cpu'''

              SCREAMING_SNAKE_CASE_							:				Dict		=self.get_dummy_components()
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.pipeline_class(**_A				)
              pipe.to(_A				)
              pipe.set_progress_bar_config(disable=_A				)

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=self.get_dummy_inputs(_A				)
              SCREAMING_SNAKE_CASE_							:				str		=pipe(**_A				).images
              SCREAMING_SNAKE_CASE_							:				Tuple		=image[0, -3:, -3:, -1]

              self.assertEqual(image.shape						,       (1, 256, 256, 3)				)
              SCREAMING_SNAKE_CASE_							:				List[str]		=np.array(
                  [0.47_222_412, 0.41_921_633, 0.44_717_434, 0.46_874_192, 0.42_588_258, 0.46_150_726, 0.4_677_534, 0.45_583_832, 0.48_579_055]				)
              SCREAMING_SNAKE_CASE_							:				Any		=np.abs(image_slice.flatten() - expected_slice				).max()
              self.assertLessEqual(_A						,       1e-3				)


       def 	_snake_case				(    self				)						->    Tuple:
              super().test_attention_slicing_forward_pass(expected_max_diff=7e-3				)


       def 	_snake_case				(    self				)						->    List[Any]:
              super().test_cpu_offload_forward_pass(expected_max_diff=3e-3				)


       def 	_snake_case				(    self				)						->    Optional[Any]:
              super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3				)


       def 	_snake_case				(    self				)						->    Tuple:
              super().test_inference_batch_single_identical(expected_max_diff=7e-3				)


       def 	_snake_case				(    self				)						->    List[Any]:
              super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3				)


       def 	_snake_case				(    self				)						->    Optional[int]:
              super().test_save_load_local(expected_max_difference=3e-3				)


       def 	_snake_case				(    self				)						->    Optional[Any]:
              super().test_save_load_optional_components(expected_max_difference=3e-3				)





       def 	_snake_case				(    self				)						->    Optional[int]:
              SCREAMING_SNAKE_CASE_							:				List[str]		=[
                  '''DDIMScheduler''',
                  '''DDPMScheduler''',
                  '''PNDMScheduler''',
                  '''HeunDiscreteScheduler''',
                  '''EulerAncestralDiscreteScheduler''',
                  '''KDPM2DiscreteScheduler''',
                  '''KDPM2AncestralDiscreteScheduler''',
                  '''DPMSolverSDEScheduler''',
              ]
              SCREAMING_SNAKE_CASE_							:				List[Any]		=self.get_dummy_components()
              SCREAMING_SNAKE_CASE_							:				Dict		=self.pipeline_class(**_A				)

              # make sure that PNDM does not need warm-up
              pipe.scheduler.register_to_config(skip_prk_steps=_A				)

              pipe.to(_A				)
              pipe.set_progress_bar_config(disable=_A				)
              SCREAMING_SNAKE_CASE_							:				Dict		=self.get_dummy_inputs(_A				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=2

              SCREAMING_SNAKE_CASE_							:				int		=[]
              for scheduler_enum in KarrasDiffusionSchedulers:
                     if scheduler_enum.name in skip_schedulers:
                            # no sigma schedulers are not supported
                            # no schedulers
                            continue

                     SCREAMING_SNAKE_CASE_							:				List[Any]		=getattr(_A						,       scheduler_enum.name				)
                     SCREAMING_SNAKE_CASE_							:				Any		=scheduler_cls.from_config(pipe.scheduler.config				)
                     SCREAMING_SNAKE_CASE_							:				Optional[int]		=pipe(**_A				)[0]
                     outputs.append(_A				)

              assert check_same_shape(_A				)






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


       def 	_snake_case				(    self				)						->    Optional[int]:
              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()


       def 	_snake_case				(    self				)						->    Any:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=torch.manual_seed(33				)

              SCREAMING_SNAKE_CASE_							:				List[Any]		=StableDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''						,       torch_dtype=torch.floataa				)
              pipe.to('''cuda'''				)

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=StableDiffusionLatentUpscalePipeline.from_pretrained(
                  '''stabilityai/sd-x2-latent-upscaler'''						,       torch_dtype=torch.floataa				)
              upscaler.to('''cuda'''				)

              SCREAMING_SNAKE_CASE_							:				str		='''a photo of an astronaut high resolution, unreal engine, ultra realistic'''

              SCREAMING_SNAKE_CASE_							:				Tuple		=pipe(_A						,       generator=_A						,       output_type='''latent'''				).images

              SCREAMING_SNAKE_CASE_							:				Tuple		=upscaler(
                  prompt=_A						,       image=_A						,       num_inference_steps=20						,       guidance_scale=0						,       generator=_A						,       output_type='''np'''						,       ).images[0]

              SCREAMING_SNAKE_CASE_							:				str		=load_numpy(
                  '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'''				)
              assert np.abs((expected_image - image).mean()				) < 5e-2





       def 	_snake_case				(    self				)						->    Dict:
              SCREAMING_SNAKE_CASE_							:				str		=torch.manual_seed(33				)

              SCREAMING_SNAKE_CASE_							:				Dict		=StableDiffusionLatentUpscalePipeline.from_pretrained(
                  '''stabilityai/sd-x2-latent-upscaler'''						,       torch_dtype=torch.floataa				)
              upscaler.to('''cuda'''				)

              SCREAMING_SNAKE_CASE_							:				List[str]		='''the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'''

              SCREAMING_SNAKE_CASE_							:				Dict		=load_image(
                  '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'''				)

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=upscaler(
                  prompt=_A						,       image=_A						,       num_inference_steps=20						,       guidance_scale=0						,       generator=_A						,       output_type='''np'''						,       ).images[0]

              SCREAMING_SNAKE_CASE_							:				int		=load_numpy(
                  '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'''				)
              assert np.abs((expected_image - image).max()				) < 5e-2

from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
					import tensorflow as tf

					from transformers import TFBlipTextModel
					from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST
class     lowerCAmelCase_     :







							def __init__(   self      :     Any ,			_A      :     int ,			_A      :     int=12 ,			_A      :     int=7 ,			_A      :     Tuple=True ,			_A      :     Optional[int]=True ,			_A      :     Union[str, Any]=True ,			_A      :     str=99 ,			_A      :     str=32 ,			_A      :     int=32 ,			_A      :     Optional[Any]=2 ,			_A      :     Dict=4 ,			_A      :     int=37 ,			_A      :     List[Any]=0.1 ,			_A      :     str=0.1 ,			_A      :     Any=512 ,			_A      :     int=0.02 ,			_A      :     Optional[Any]=0 ,			_A      :     Dict=None ,			):
												_UpperCamelCase			     =     parent
												_UpperCamelCase			     =     batch_size
												_UpperCamelCase			     =     seq_length
												_UpperCamelCase			     =     is_training
												_UpperCamelCase			     =     use_input_mask
												_UpperCamelCase			     =     use_labels
												_UpperCamelCase			     =     vocab_size
												_UpperCamelCase			     =     hidden_size
												_UpperCamelCase			     =     projection_dim
												_UpperCamelCase			     =     num_hidden_layers
												_UpperCamelCase			     =     num_attention_heads
												_UpperCamelCase			     =     intermediate_size
												_UpperCamelCase			     =     dropout
												_UpperCamelCase			     =     attention_dropout
												_UpperCamelCase			     =     max_position_embeddings
												_UpperCamelCase			     =     initializer_range
												_UpperCamelCase			     =     scope
												_UpperCamelCase			     =     bos_token_id







							def  UpperCamelCase_						(   self      :     List[str]		):
												_UpperCamelCase			     =     ids_tensor([self.batch_size, self.seq_length] ,			self.vocab_size		)

												_UpperCamelCase			     =     None
												if self.use_input_mask:
																	_UpperCamelCase			     =     random_attention_mask([self.batch_size, self.seq_length]		)

												if input_mask is not None:
																	_UpperCamelCase			     =     input_mask.numpy()
																	_UpperCamelCase     ,    _UpperCamelCase			     =     input_mask.shape
																	_UpperCamelCase			     =     np.random.randint(1 ,			seq_length - 1 ,			size=(batch_size,)		)
																	for batch_idx, start_index in enumerate(_A		):
																						_UpperCamelCase			     =     1
																						_UpperCamelCase			     =     0

												_UpperCamelCase			     =     self.get_config()

												return config, input_ids, tf.convert_to_tensor(_A		)







							def  UpperCamelCase_						(   self      :     str		):
												return BlipTextConfig(
												    vocab_size=self.vocab_size ,			hidden_size=self.hidden_size ,			projection_dim=self.projection_dim ,			num_hidden_layers=self.num_hidden_layers ,			num_attention_heads=self.num_attention_heads ,			intermediate_size=self.intermediate_size ,			dropout=self.dropout ,			attention_dropout=self.attention_dropout ,			max_position_embeddings=self.max_position_embeddings ,			initializer_range=self.initializer_range ,			bos_token_id=self.bos_token_id ,			)







							def  UpperCamelCase_						(   self      :     List[str] ,			_A      :     Tuple ,			_A      :     str ,			_A      :     Optional[Any]		):
												_UpperCamelCase			     =     TFBlipTextModel(config=_A		)
												_UpperCamelCase			     =     model(_A ,			attention_mask=_A ,			training=_A		)
												_UpperCamelCase			     =     model(_A ,			training=_A		)
												self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size)		)
												self.parent.assertEqual(result.pooler_output.shape ,			(self.batch_size, self.hidden_size)		)







							def  UpperCamelCase_						(   self      :     Tuple		):
												_UpperCamelCase			     =     self.prepare_config_and_inputs()
												_UpperCamelCase     ,    _UpperCamelCase     ,    _UpperCamelCase			     =     config_and_inputs
												_UpperCamelCase			     =     {'''input_ids''': input_ids, '''attention_mask''': input_mask}
												return config, inputs_dict




@require_tf
class     lowerCAmelCase_     (       __lowercase,  unittest.TestCase ):
							UpperCAmelCase				   =     (TFBlipTextModel,) if is_tf_available() else ()
							UpperCAmelCase				   =     False
							UpperCAmelCase				   =     False
							UpperCAmelCase				   =     False







							def  UpperCamelCase_						(   self      :     Dict		):
												_UpperCamelCase			     =     BlipTextModelTester(self		)
												_UpperCamelCase			     =     ConfigTester(self ,			config_class=_A ,			hidden_size=37		)







							def  UpperCamelCase_						(   self      :     Dict		):
												self.config_tester.run_common_tests()







							def  UpperCamelCase_						(   self      :     int		):
												_UpperCamelCase			     =     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_A		)







							def  UpperCamelCase_						(   self      :     List[Any]		):
												pass







							def  UpperCamelCase_						(   self      :     Tuple		):
												pass







							@unittest.skip(reason='''Blip does not use inputs_embeds'''		)
							def  UpperCamelCase_						(   self      :     Dict		):
												pass







							@unittest.skip(reason='''BlipTextModel has no base class and is not available in MODEL_MAPPING'''		)
							def  UpperCamelCase_						(   self      :     Dict		):
												pass







							@unittest.skip(reason='''BlipTextModel has no base class and is not available in MODEL_MAPPING'''		)
							def  UpperCamelCase_						(   self      :     List[str]		):
												pass







							@slow
							def  UpperCamelCase_						(   self      :     Optional[int]		):
												for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_UpperCamelCase			     =     TFBlipTextModel.from_pretrained(_A		)
																	self.assertIsNotNone(_A		)







							def  UpperCamelCase_						(   self      :     int ,			_A      :     Optional[int]=True		):
												super().test_pt_tf_model_equivalence(allow_missing_keys=_A		)
