

'''simple docstring'''
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class 					_snake_case		(lowerCAmelCase__			,			lowerCAmelCase__			,			unittest.TestCase):
			__A  :	str        =StableDiffusionSAGPipeline
			__A  :	str        =TEXT_TO_IMAGE_PARAMS
			__A  :	int        =TEXT_TO_IMAGE_BATCH_PARAMS
			__A  :	Tuple        =TEXT_TO_IMAGE_IMAGE_PARAMS
			__A  :	Dict        =TEXT_TO_IMAGE_IMAGE_PARAMS
			__A  :	List[Any]        =False






			def 		UpperCamelCase__							(  self     ):
				torch.manual_seed(0     )
				UpperCAmelCase_     :      List[Any]						=  UNetaDConditionModel(
				    block_out_channels=(32, 64)			,layers_per_block=2			,sample_size=32			,in_channels=4			,out_channels=4			,down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")			,up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")			,cross_attention_dim=32			,)
				UpperCAmelCase_     :      Dict						=  DDIMScheduler(
				    beta_start=0.00085			,beta_end=0.012			,beta_schedule="scaled_linear"			,clip_sample=__UpperCAmelCase			,set_alpha_to_one=__UpperCAmelCase			,)
				torch.manual_seed(0     )
				UpperCAmelCase_     :      List[str]						=  AutoencoderKL(
				    block_out_channels=[32, 64]			,in_channels=3			,out_channels=3			,down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,latent_channels=4			,)
				torch.manual_seed(0     )
				UpperCAmelCase_     :      Optional[int]						=  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			,)
				UpperCAmelCase_     :      Optional[Any]						=  CLIPTextModel(__UpperCAmelCase     )
				UpperCAmelCase_     :      int						=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )

				UpperCAmelCase_     :      Any						=  {
				    "unet": unet,
				    "scheduler": scheduler,
				    "vae": vae,
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
				return components






			def 		UpperCamelCase__							(  self			,_snake_case			,_snake_case=0     ):
				if str(__UpperCAmelCase     ).startswith("mps"     ):
					UpperCAmelCase_     :      Dict						=  torch.manual_seed(__UpperCAmelCase     )
				else:
					UpperCAmelCase_     :      List[str]						=  torch.Generator(device=__UpperCAmelCase     ).manual_seed(__UpperCAmelCase     )
				UpperCAmelCase_     :      Optional[Any]						=  {
				    "prompt": ".",
				    "generator": generator,
				    "num_inference_steps": 2,
				    "guidance_scale": 1.0,
				    "sag_scale": 1.0,
				    "output_type": "numpy",
				}
				return inputs






			def 		UpperCamelCase__							(  self     ):
				super().test_inference_batch_single_identical(expected_max_diff=3E-3     )




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






			def 		UpperCamelCase__							(  self     ):
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()






			def 		UpperCamelCase__							(  self     ):
				UpperCAmelCase_     :      List[Any]						=  StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"     )
				UpperCAmelCase_     :      Union[str, Any]						=  sag_pipe.to(__UpperCAmelCase     )
				sag_pipe.set_progress_bar_config(disable=__UpperCAmelCase     )

				UpperCAmelCase_     :      Optional[int]						=  "."
				UpperCAmelCase_     :      Any						=  torch.manual_seed(0     )
				UpperCAmelCase_     :      Optional[Any]						=  sag_pipe(
				    [prompt]			,generator=__UpperCAmelCase			,guidance_scale=7.5			,sag_scale=1.0			,num_inference_steps=20			,output_type="np"     )

				UpperCAmelCase_     :      int						=  output.images

				UpperCAmelCase_     :      Union[str, Any]						=  image[0, -3:, -3:, -1]

				assert image.shape == (1, 5_12, 5_12, 3)
				UpperCAmelCase_     :      Any						=  np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]     )

				assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5E-2






			def 		UpperCamelCase__							(  self     ):
				UpperCAmelCase_     :      List[Any]						=  StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"     )
				UpperCAmelCase_     :      List[Any]						=  sag_pipe.to(__UpperCAmelCase     )
				sag_pipe.set_progress_bar_config(disable=__UpperCAmelCase     )

				UpperCAmelCase_     :      Any						=  "."
				UpperCAmelCase_     :      List[Any]						=  torch.manual_seed(0     )
				UpperCAmelCase_     :      List[str]						=  sag_pipe(
				    [prompt]			,generator=__UpperCAmelCase			,guidance_scale=7.5			,sag_scale=1.0			,num_inference_steps=20			,output_type="np"     )

				UpperCAmelCase_     :      Dict						=  output.images

				UpperCAmelCase_     :      Dict						=  image[0, -3:, -3:, -1]

				assert image.shape == (1, 5_12, 5_12, 3)
				UpperCAmelCase_     :      Optional[int]						=  np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]     )

				assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5E-2






			def 		UpperCamelCase__							(  self     ):
				UpperCAmelCase_     :      str						=  StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"     )
				UpperCAmelCase_     :      List[Any]						=  sag_pipe.to(__UpperCAmelCase     )
				sag_pipe.set_progress_bar_config(disable=__UpperCAmelCase     )

				UpperCAmelCase_     :      str						=  "."
				UpperCAmelCase_     :      Dict						=  torch.manual_seed(0     )
				UpperCAmelCase_     :      Optional[Any]						=  sag_pipe(
				    [prompt]			,width=7_68			,height=5_12			,generator=__UpperCAmelCase			,guidance_scale=7.5			,sag_scale=1.0			,num_inference_steps=20			,output_type="np"			,)

				UpperCAmelCase_     :      Any						=  output.images

				assert image.shape == (1, 5_12, 7_68, 3)









import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      SqueezeBertForMaskedLM,
      SqueezeBertForMultipleChoice,
      SqueezeBertForQuestionAnswering,
      SqueezeBertForSequenceClassification,
      SqueezeBertForTokenClassification,
      SqueezeBertModel,
  )



class 			__lowerCAmelCase  (						lowerCAmelCase__ ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase=13						,					__UpperCAmelCase=7						,					__UpperCAmelCase=True						,					__UpperCAmelCase=True						,					__UpperCAmelCase=False						,					__UpperCAmelCase=True						,					__UpperCAmelCase=99						,					__UpperCAmelCase=32						,					__UpperCAmelCase=5						,					__UpperCAmelCase=4						,					__UpperCAmelCase=64						,					__UpperCAmelCase="gelu"						,					__UpperCAmelCase=0.1						,					__UpperCAmelCase=0.1						,					__UpperCAmelCase=512						,					__UpperCAmelCase=16						,					__UpperCAmelCase=2						,					__UpperCAmelCase=0.02						,					__UpperCAmelCase=3						,					__UpperCAmelCase=4						,					__UpperCAmelCase=None						,					__UpperCAmelCase=2						,					__UpperCAmelCase=2						,					__UpperCAmelCase=2						,					__UpperCAmelCase=2						,					__UpperCAmelCase=4						,					__UpperCAmelCase=1						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     parent
    __lowerCamelCase    				=     batch_size
    __lowerCamelCase    				=     seq_length
    __lowerCamelCase    				=     is_training
    __lowerCamelCase    				=     use_input_mask
    __lowerCamelCase    				=     use_token_type_ids
    __lowerCamelCase    				=     use_labels
    __lowerCamelCase    				=     vocab_size
    __lowerCamelCase    				=     hidden_size
    __lowerCamelCase    				=     num_hidden_layers
    __lowerCamelCase    				=     num_attention_heads
    __lowerCamelCase    				=     intermediate_size
    __lowerCamelCase    				=     hidden_act
    __lowerCamelCase    				=     hidden_dropout_prob
    __lowerCamelCase    				=     attention_probs_dropout_prob
    __lowerCamelCase    				=     max_position_embeddings
    __lowerCamelCase    				=     type_vocab_size
    __lowerCamelCase    				=     type_sequence_label_size
    __lowerCamelCase    				=     initializer_range
    __lowerCamelCase    				=     num_labels
    __lowerCamelCase    				=     num_choices
    __lowerCamelCase    				=     scope
    __lowerCamelCase    				=     q_groups
    __lowerCamelCase    				=     k_groups
    __lowerCamelCase    				=     v_groups
    __lowerCamelCase    				=     post_attention_groups
    __lowerCamelCase    				=     intermediate_groups
    __lowerCamelCase    				=     output_groups




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     ids_tensor([self.batch_size, self.seq_length]						,					self.vocab_size					)

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

    __lowerCamelCase    				=     None
    __lowerCamelCase    				=     None
    __lowerCamelCase    				=     None
    if self.use_labels:
       __lowerCamelCase    				=     ids_tensor([self.batch_size]						,					self.type_sequence_label_size					)
       __lowerCamelCase    				=     ids_tensor([self.batch_size, self.seq_length]						,					self.num_labels					)
       __lowerCamelCase    				=     ids_tensor([self.batch_size]						,					self.num_choices					)

    __lowerCamelCase    				=     self.get_config()

    return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    return SqueezeBertConfig(
        embedding_size=self.hidden_size						,					vocab_size=self.vocab_size						,					hidden_size=self.hidden_size						,					num_hidden_layers=self.num_hidden_layers						,					num_attention_heads=self.num_attention_heads						,					intermediate_size=self.intermediate_size						,					hidden_act=self.hidden_act						,					attention_probs_dropout_prob=self.hidden_dropout_prob						,					attention_dropout=self.attention_probs_dropout_prob						,					max_position_embeddings=self.max_position_embeddings						,					initializer_range=self.initializer_range						,					q_groups=self.q_groups						,					k_groups=self.k_groups						,					v_groups=self.v_groups						,					post_attention_groups=self.post_attention_groups						,					intermediate_groups=self.intermediate_groups						,					output_groups=self.output_groups						,					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertModel(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     model(__UpperCAmelCase					)
    self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, self.seq_length, self.hidden_size)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertForMaskedLM(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.seq_length, self.vocab_size)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertForQuestionAnswering(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(
        __UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					start_positions=__UpperCAmelCase						,					end_positions=__UpperCAmelCase					)
    self.parent.assertEqual(result.start_logits.shape						,					(self.batch_size, self.seq_length)					)
    self.parent.assertEqual(result.end_logits.shape						,					(self.batch_size, self.seq_length)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.num_labels
    __lowerCamelCase    				=     SqueezeBertForSequenceClassification(__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.num_labels)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.num_labels
    __lowerCamelCase    				=     SqueezeBertForTokenClassification(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()

    __lowerCamelCase    				=     model(__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.seq_length, self.num_labels)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.num_choices
    __lowerCamelCase    				=     SqueezeBertForMultipleChoice(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     input_ids.unsqueeze(1					).expand(-1						,					self.num_choices						,					-1					).contiguous()
    __lowerCamelCase    				=     input_mask.unsqueeze(1					).expand(-1						,					self.num_choices						,					-1					).contiguous()
    __lowerCamelCase    				=     model(
        __UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase						,					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.num_choices)					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.prepare_config_and_inputs()
    ((__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase))    				=     config_and_inputs
    __lowerCamelCase    				=     {'''input_ids''': input_ids, '''attention_mask''': input_mask}
    return config, inputs_dict



@require_torch
class 			__lowerCAmelCase  (						lowerCAmelCase__	,						lowerCAmelCase__	,						unittest.TestCase ):
 lowerCAmelCase__							       =	(
     (
         SqueezeBertModel,
         SqueezeBertForMaskedLM,
         SqueezeBertForMultipleChoice,
         SqueezeBertForQuestionAnswering,
         SqueezeBertForSequenceClassification,
         SqueezeBertForTokenClassification,
     )
     if is_torch_available()
     else None
 )
 lowerCAmelCase__							       =	(
     {
         """feature-extraction""": SqueezeBertModel,
         """fill-mask""": SqueezeBertForMaskedLM,
         """question-answering""": SqueezeBertForQuestionAnswering,
         """text-classification""": SqueezeBertForSequenceClassification,
         """token-classification""": SqueezeBertForTokenClassification,
         """zero-shot""": SqueezeBertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 lowerCAmelCase__							       =	False
 lowerCAmelCase__							       =	True
 lowerCAmelCase__							       =	False




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertModelTester(self					)
    __lowerCamelCase    				=     ConfigTester(self						,					config_class=__UpperCAmelCase						,					dim=37					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    self.config_tester.run_common_tests()




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_model(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_masked_lm(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_question_answering(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_sequence_classification(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_token_classification(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_multiple_choice(*__UpperCAmelCase					)




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __lowerCamelCase    				=     SqueezeBertModel.from_pretrained(__UpperCAmelCase					)
       self.assertIsNotNone(__UpperCAmelCase					)



@require_sentencepiece
@require_tokenizers
@require_torch
class 			__lowerCAmelCase  (						unittest.TestCase ):




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertForSequenceClassification.from_pretrained('''squeezebert/squeezebert-mnli'''					)

    __lowerCamelCase    				=     torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]]					)
    __lowerCamelCase    				=     model(__UpperCAmelCase					)[0]
    __lowerCamelCase    				=     torch.Size((1, 3)					)
    self.assertEqual(output.shape						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     torch.tensor([[0.6_401, -0.0_349, -0.6_041]]					)
    self.assertTrue(torch.allclose(__UpperCAmelCase						,					__UpperCAmelCase						,					atol=1E-4					)					)






