






"""simple docstring"""



import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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()







@skip_mps
class    lowercase__			(							SCREAMING_SNAKE_CASE__,   SCREAMING_SNAKE_CASE__,   unittest.TestCase   ):
			_UpperCAmelCase		:Tuple   			= StableDiffusionPanoramaPipeline
			_UpperCAmelCase		:List[str]   			= TEXT_TO_IMAGE_PARAMS
			_UpperCAmelCase		:Any   			= TEXT_TO_IMAGE_BATCH_PARAMS
			_UpperCAmelCase		:Union[str, Any]   			= TEXT_TO_IMAGE_IMAGE_PARAMS
			_UpperCAmelCase		:Tuple   			= TEXT_TO_IMAGE_IMAGE_PARAMS


			def 	UpperCAmelCase__						(      self     :  int					):
										torch.manual_seed(0					)
										lowerCamelCase_       :							str									=UNetaDConditionModel(
										    block_out_channels=(32, 64)   ,       layers_per_block=1   ,       sample_size=32   ,       in_channels=4   ,       out_channels=4   ,       down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")   ,       up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")   ,       cross_attention_dim=32   ,       )
										lowerCamelCase_       :							str									=DDIMScheduler()
										torch.manual_seed(0					)
										lowerCamelCase_       :							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					)
										lowerCamelCase_       :							Tuple									=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=1000   ,       )
										lowerCamelCase_       :							List[Any]									=CLIPTextModel(lowerCamelCase__					)
										lowerCamelCase_       :							Tuple									=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

										lowerCamelCase_       :							List[str]									={
										    """unet""": unet,
										    """scheduler""": scheduler,
										    """vae""": vae,
										    """text_encoder""": text_encoder,
										    """tokenizer""": tokenizer,
										    """safety_checker""": None,
										    """feature_extractor""": None,
										}
										return components


			def 	UpperCAmelCase__						(      self     :  Optional[int]   ,       snake_case__     :  str   ,       snake_case__     :  List[str]=0					):
										lowerCamelCase_       :							List[Any]									=torch.manual_seed(lowerCamelCase__					)
										lowerCamelCase_       :							Tuple									={
										    """prompt""": """a photo of the dolomites""",
										    """generator""": generator,
										    # Setting height and width to None to prevent OOMs on CPU.
										    """height""": None,
										    """width""": None,
										    """num_inference_steps""": 1,
										    """guidance_scale""": 6.0,
										    """output_type""": """numpy""",
										}
										return inputs


			def 	UpperCAmelCase__						(      self     :  Tuple					):
										lowerCamelCase_       :							Optional[int]									="""cpu"""  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_       :							Any									=self.get_dummy_components()
										lowerCamelCase_       :							Optional[Any]									=StableDiffusionPanoramaPipeline(**lowerCamelCase__					)
										lowerCamelCase_       :							Optional[int]									=sd_pipe.to(lowerCamelCase__					)
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase__					)

										lowerCamelCase_       :							List[Any]									=self.get_dummy_inputs(lowerCamelCase__					)
										lowerCamelCase_       :							Tuple									=sd_pipe(**lowerCamelCase__					).images
										lowerCamelCase_       :							Optional[int]									=image[0, -3:, -3:, -1]
										assert image.shape == (1, 64, 64, 3)

										lowerCamelCase_       :							Optional[Any]									=np.array([0.6_186, 0.5_374, 0.4_915, 0.4_135, 0.4_114, 0.4_563, 0.5_128, 0.4_977, 0.4_757]					)

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


			def 	UpperCAmelCase__						(      self     :  int					):
										super().test_inference_batch_consistent(batch_sizes=[1, 2]					)


			def 	UpperCAmelCase__						(      self     :  Dict					):
										super().test_inference_batch_single_identical(batch_size=2   ,       expected_max_diff=3.25E-3					)


			def 	UpperCAmelCase__						(      self     :  int					):
										lowerCamelCase_       :							Tuple									="""cpu"""  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_       :							Optional[Any]									=self.get_dummy_components()
										lowerCamelCase_       :							Optional[Any]									=StableDiffusionPanoramaPipeline(**lowerCamelCase__					)
										lowerCamelCase_       :							List[str]									=sd_pipe.to(lowerCamelCase__					)
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase__					)

										lowerCamelCase_       :							Optional[int]									=self.get_dummy_inputs(lowerCamelCase__					)
										lowerCamelCase_       :							Any									="""french fries"""
										lowerCamelCase_       :							int									=sd_pipe(**lowerCamelCase__   ,       negative_prompt=lowerCamelCase__					)
										lowerCamelCase_       :							Dict									=output.images
										lowerCamelCase_       :							str									=image[0, -3:, -3:, -1]

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

										lowerCamelCase_       :							Tuple									=np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]					)

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


			def 	UpperCAmelCase__						(      self     :  Dict					):
										lowerCamelCase_       :							Optional[int]									="""cpu"""  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_       :							str									=self.get_dummy_components()
										lowerCamelCase_       :							int									=StableDiffusionPanoramaPipeline(**lowerCamelCase__					)
										lowerCamelCase_       :							Optional[Any]									=sd_pipe.to(lowerCamelCase__					)
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase__					)

										lowerCamelCase_       :							Any									=self.get_dummy_inputs(lowerCamelCase__					)
										lowerCamelCase_       :							Optional[Any]									=sd_pipe(**lowerCamelCase__   ,       view_batch_size=2					)
										lowerCamelCase_       :							Optional[int]									=output.images
										lowerCamelCase_       :							Optional[Any]									=image[0, -3:, -3:, -1]

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

										lowerCamelCase_       :							List[str]									=np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]					)

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


			def 	UpperCAmelCase__						(      self     :  Tuple					):
										lowerCamelCase_       :							List[str]									="""cpu"""  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_       :							List[str]									=self.get_dummy_components()
										lowerCamelCase_       :							Tuple									=EulerAncestralDiscreteScheduler(
										    beta_start=0.00_085   ,       beta_end=0.012   ,       beta_schedule="scaled_linear"					)
										lowerCamelCase_       :							int									=StableDiffusionPanoramaPipeline(**lowerCamelCase__					)
										lowerCamelCase_       :							List[str]									=sd_pipe.to(lowerCamelCase__					)
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase__					)

										lowerCamelCase_       :							Tuple									=self.get_dummy_inputs(lowerCamelCase__					)
										lowerCamelCase_       :							int									=sd_pipe(**lowerCamelCase__					).images
										lowerCamelCase_       :							Any									=image[0, -3:, -3:, -1]

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

										lowerCamelCase_       :							Tuple									=np.array([0.4_024, 0.6_510, 0.4_901, 0.5_378, 0.5_813, 0.5_622, 0.4_795, 0.4_467, 0.4_952]					)

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


			def 	UpperCAmelCase__						(      self     :  Optional[int]					):
										lowerCamelCase_       :							List[Any]									="""cpu"""  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_       :							List[Any]									=self.get_dummy_components()
										lowerCamelCase_       :							List[str]									=PNDMScheduler(
										    beta_start=0.00_085   ,       beta_end=0.012   ,       beta_schedule="scaled_linear"   ,       skip_prk_steps=lowerCamelCase__					)
										lowerCamelCase_       :							Tuple									=StableDiffusionPanoramaPipeline(**lowerCamelCase__					)
										lowerCamelCase_       :							List[str]									=sd_pipe.to(lowerCamelCase__					)
										sd_pipe.set_progress_bar_config(disable=lowerCamelCase__					)

										lowerCamelCase_       :							List[Any]									=self.get_dummy_inputs(lowerCamelCase__					)
										lowerCamelCase_       :							int									=sd_pipe(**lowerCamelCase__					).images
										lowerCamelCase_       :							Union[str, Any]									=image[0, -3:, -3:, -1]

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

										lowerCamelCase_       :							Optional[int]									=np.array([0.6_391, 0.6_291, 0.4_861, 0.5_134, 0.5_552, 0.4_578, 0.5_032, 0.5_023, 0.4_539]					)

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







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


			def 	UpperCAmelCase__						(      self     :  List[str]					):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


			def 	UpperCAmelCase__						(      self     :  Optional[Any]   ,       snake_case__     :  Dict=0					):
										lowerCamelCase_       :							Union[str, Any]									=torch.manual_seed(lowerCamelCase__					)
										lowerCamelCase_       :							Optional[Any]									={
										    """prompt""": """a photo of the dolomites""",
										    """generator""": generator,
										    """num_inference_steps""": 3,
										    """guidance_scale""": 7.5,
										    """output_type""": """numpy""",
										}
										return inputs


			def 	UpperCAmelCase__						(      self     :  Optional[int]					):
										lowerCamelCase_       :							str									="""stabilityai/stable-diffusion-2-base"""
										lowerCamelCase_       :							Dict									=DDIMScheduler.from_pretrained(lowerCamelCase__   ,       subfolder="scheduler"					)
										lowerCamelCase_       :							int									=StableDiffusionPanoramaPipeline.from_pretrained(lowerCamelCase__   ,       scheduler=lowerCamelCase__   ,       safety_checker=lowerCamelCase__					)
										pipe.to(lowerCamelCase__					)
										pipe.set_progress_bar_config(disable=lowerCamelCase__					)
										pipe.enable_attention_slicing()

										lowerCamelCase_       :							Optional[int]									=self.get_inputs()
										lowerCamelCase_       :							List[str]									=pipe(**lowerCamelCase__					).images
										lowerCamelCase_       :							Union[str, Any]									=image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 512, 2048, 3)

										lowerCamelCase_       :							str									=np.array(
										    [
										        0.36_968_392,
										        0.27_025_372,
										        0.32_446_766,
										        0.28_379_387,
										        0.36_363_274,
										        0.30_733_347,
										        0.27_100_027,
										        0.27_054_125,
										        0.25_536_096,
										    ]					)

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


			def 	UpperCAmelCase__						(      self     :  Optional[int]					):
										lowerCamelCase_       :							Optional[Any]									=StableDiffusionPanoramaPipeline.from_pretrained(
										    "stabilityai/stable-diffusion-2-base"   ,       safety_checker=lowerCamelCase__					)
										lowerCamelCase_       :							List[Any]									=LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
										pipe.to(lowerCamelCase__					)
										pipe.set_progress_bar_config(disable=lowerCamelCase__					)
										pipe.enable_attention_slicing()

										lowerCamelCase_       :							Optional[Any]									=self.get_inputs()
										lowerCamelCase_       :							int									=pipe(**lowerCamelCase__					).images
										lowerCamelCase_       :							List[str]									=image[0, -3:, -3:, -1].flatten()

										assert image.shape == (1, 512, 2048, 3)

										lowerCamelCase_       :							List[Any]									=np.array(
										    [
										        [
										            0.0,
										            0.0,
										            0.0,
										            0.0,
										            0.0,
										            0.0,
										            0.0,
										            0.0,
										            0.0,
										        ]
										    ]					)

										assert np.abs(expected_slice - image_slice					).max() < 1E-3


			def 	UpperCAmelCase__						(      self     :  List[str]					):
										lowerCamelCase_       :							Dict									=0

										def callback_fn(snake_case__     :  int   ,       snake_case__     :  Optional[Any]   ,       snake_case__     :  List[str]					) -> None:
																	lowerCamelCase_       :							int									=True
																	nonlocal number_of_steps
																	number_of_steps += 1
																	if step == 1:
																								lowerCamelCase_       :							Dict									=latents.detach().cpu().numpy()
																								assert latents.shape == (1, 4, 64, 256)
																								lowerCamelCase_       :							Optional[int]									=latents[0, -3:, -3:, -1]

																								lowerCamelCase_       :							str									=np.array(
																								    [
																								        0.18_681_869,
																								        0.33_907_816,
																								        0.5_361_276,
																								        0.14_432_865,
																								        -0.02_856_611,
																								        -0.73_941_123,
																								        0.23_397_987,
																								        0.47_322_682,
																								        -0.37_823_164,
																								    ]					)
																								assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2
																	elif step == 2:
																								lowerCamelCase_       :							Any									=latents.detach().cpu().numpy()
																								assert latents.shape == (1, 4, 64, 256)
																								lowerCamelCase_       :							List[Any]									=latents[0, -3:, -3:, -1]

																								lowerCamelCase_       :							int									=np.array(
																								    [
																								        0.18_539_645,
																								        0.33_987_248,
																								        0.5_378_559,
																								        0.14_437_142,
																								        -0.02_455_261,
																								        -0.7_338_317,
																								        0.23_990_755,
																								        0.47_356_272,
																								        -0.3_786_505,
																								    ]					)

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

										lowerCamelCase_       :							List[str]									=False

										lowerCamelCase_       :							int									="""stabilityai/stable-diffusion-2-base"""
										lowerCamelCase_       :							Dict									=DDIMScheduler.from_pretrained(lowerCamelCase__   ,       subfolder="scheduler"					)
										lowerCamelCase_       :							int									=StableDiffusionPanoramaPipeline.from_pretrained(lowerCamelCase__   ,       scheduler=lowerCamelCase__   ,       safety_checker=lowerCamelCase__					)
										lowerCamelCase_       :							Dict									=pipe.to(lowerCamelCase__					)
										pipe.set_progress_bar_config(disable=lowerCamelCase__					)
										pipe.enable_attention_slicing()

										lowerCamelCase_       :							Union[str, Any]									=self.get_inputs()
										pipe(**lowerCamelCase__   ,       callback=lowerCamelCase__   ,       callback_steps=1					)
										assert callback_fn.has_been_called
										assert number_of_steps == 3


			def 	UpperCAmelCase__						(      self     :  Optional[Any]					):
										torch.cuda.empty_cache()
										torch.cuda.reset_max_memory_allocated()
										torch.cuda.reset_peak_memory_stats()

										lowerCamelCase_       :							List[str]									="""stabilityai/stable-diffusion-2-base"""
										lowerCamelCase_       :							Optional[int]									=DDIMScheduler.from_pretrained(lowerCamelCase__   ,       subfolder="scheduler"					)
										lowerCamelCase_       :							Tuple									=StableDiffusionPanoramaPipeline.from_pretrained(lowerCamelCase__   ,       scheduler=lowerCamelCase__   ,       safety_checker=lowerCamelCase__					)
										lowerCamelCase_       :							List[Any]									=pipe.to(lowerCamelCase__					)
										pipe.set_progress_bar_config(disable=lowerCamelCase__					)
										pipe.enable_attention_slicing(1					)
										pipe.enable_sequential_cpu_offload()

										lowerCamelCase_       :							Optional[Any]									=self.get_inputs()
										lowerCamelCase_       :							Union[str, Any]									=pipe(**lowerCamelCase__					)

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








import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
     import jax.numpy as jnp

     from transformers.models.albert.modeling_flax_albert import (
         FlaxAlbertForMaskedLM,
         FlaxAlbertForMultipleChoice,
         FlaxAlbertForPreTraining,
         FlaxAlbertForQuestionAnswering,
         FlaxAlbertForSequenceClassification,
         FlaxAlbertForTokenClassification,
         FlaxAlbertModel,
     )






class SCREAMING_SNAKE_CASE__     (		unittest.TestCase    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__=13,		lowerCamelCase__=7,		lowerCamelCase__=True,		lowerCamelCase__=True,		lowerCamelCase__=True,		lowerCamelCase__=True,		lowerCamelCase__=99,		lowerCamelCase__=32,		lowerCamelCase__=5,		lowerCamelCase__=4,		lowerCamelCase__=37,		lowerCamelCase__="gelu",		lowerCamelCase__=0.1,		lowerCamelCase__=0.1,		lowerCamelCase__=512,		lowerCamelCase__=16,		lowerCamelCase__=2,		lowerCamelCase__=0.02,		lowerCamelCase__=4,		):
        A :       List[str]          =   parent
        A :       Optional[int]          =   batch_size
        A :       Union[str, Any]          =   seq_length
        A :       Any          =   is_training
        A :       List[str]          =   use_attention_mask
        A :       Union[str, Any]          =   use_token_type_ids
        A :       Any          =   use_labels
        A :       str          =   vocab_size
        A :       Union[str, Any]          =   hidden_size
        A :       str          =   num_hidden_layers
        A :       List[Any]          =   num_attention_heads
        A :       Optional[int]          =   intermediate_size
        A :       Optional[Any]          =   hidden_act
        A :       Dict          =   hidden_dropout_prob
        A :       List[Any]          =   attention_probs_dropout_prob
        A :       Optional[int]          =   max_position_embeddings
        A :       int          =   type_vocab_size
        A :       str          =   type_sequence_label_size
        A :       List[Any]          =   initializer_range
        A :       str          =   num_choices






    def 			_lowerCAmelCase					(       self		):
        A :       Optional[int]          =   ids_tensor([self.batch_size, self.seq_length],		self.vocab_size		)

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

        A :       int          =   None
        if self.use_token_type_ids:
            A :       List[Any]          =   ids_tensor([self.batch_size, self.seq_length],		self.type_vocab_size		)

        A :       Optional[int]          =   AlbertConfig(
            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,		hidden_dropout_prob=self.hidden_dropout_prob,		attention_probs_dropout_prob=self.attention_probs_dropout_prob,		max_position_embeddings=self.max_position_embeddings,		type_vocab_size=self.type_vocab_size,		is_decoder=lowerCamelCase__,		initializer_range=self.initializer_range,		)

        return config, input_ids, token_type_ids, attention_mask






    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   self.prepare_config_and_inputs()
        A     ,							A     ,							A     ,							A :       str          =   config_and_inputs
        A :       Union[str, Any]          =   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
        return config, inputs_dict






@require_flax
class SCREAMING_SNAKE_CASE__     (		SCREAMING_SNAKE_CASE__ ,						unittest.TestCase    ):




    '''simple docstring'''






    __lowerCamelCase			: Any					=      (
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )






    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   FlaxAlbertModelTester(self		)






    @slow
    def 			_lowerCAmelCase					(       self		):
        for model_class_name in self.all_model_classes:
            A :       Dict          =   model_class_name.from_pretrained("""albert-base-v2"""		)
            A :       Optional[int]          =   model(np.ones((1, 1)		)		)
            self.assertIsNotNone(lowerCamelCase__		)






@require_flax
class SCREAMING_SNAKE_CASE__     (		unittest.TestCase    ):




    '''simple docstring'''






    @slow
    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   FlaxAlbertModel.from_pretrained("""albert-base-v2"""		)
        A :       List[str]          =   np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]		)
        A :       str          =   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
        A :       Optional[int]          =   model(lowerCamelCase__,		attention_mask=lowerCamelCase__		)[0]
        A :       str          =   (1, 11, 768)
        self.assertEqual(output.shape,		lowerCamelCase__		)
        A :       Optional[int]          =   np.array(
            [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]		)

        self.assertTrue(jnp.allclose(output[:, 1:4, 1:4],		lowerCamelCase__,		atol=1e-4		)		)


