"""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 				_A		( lowerCAmelCase						,					lowerCAmelCase						,					unittest.TestCase			):
      snake_case__    :     int													=			StableDiffusionPanoramaPipeline
      snake_case__    :     int													=			TEXT_TO_IMAGE_PARAMS
      snake_case__    :     str													=			TEXT_TO_IMAGE_BATCH_PARAMS
      snake_case__    :     str													=			TEXT_TO_IMAGE_IMAGE_PARAMS
      snake_case__    :     Optional[Any]													=			TEXT_TO_IMAGE_IMAGE_PARAMS

      def        A__ (		self		):







            """simple docstring"""



            torch.manual_seed(0		)
            lowercase														=			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	, )
            lowercase														=			DDIMScheduler()
            torch.manual_seed(0		)
            lowercase														=			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		)
            lowercase														=			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	, )
            lowercase														=			CLIPTextModel(__lowerCAmelCase		)
            lowercase														=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)

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

      def        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase=0		):







            """simple docstring"""



            lowercase														=			torch.manual_seed(__lowerCAmelCase		)
            lowercase														=			{
                """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        A__ (		self		):







            """simple docstring"""



            lowercase														=			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowercase														=			self.get_dummy_components()
            lowercase														=			StableDiffusionPanoramaPipeline(**__lowerCAmelCase		)
            lowercase														=			sd_pipe.to(__lowerCAmelCase		)
            sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

            lowercase														=			self.get_dummy_inputs(__lowerCAmelCase		)
            lowercase														=			sd_pipe(**__lowerCAmelCase		).images
            lowercase														=			image[0, -3:, -3:, -1]
            assert image.shape == (1, 64, 64, 3)

            lowercase														=			np.array([0.6_1_8_6, 0.5_3_7_4, 0.4_9_1_5, 0.4_1_3_5, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_7, 0.4_7_5_7]		)

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

      def        A__ (		self		):







            """simple docstring"""



            super().test_inference_batch_consistent(batch_sizes=[1, 2]		)

      def        A__ (		self		):







            """simple docstring"""



            super().test_inference_batch_single_identical(batch_size=2	, expected_max_diff=3.25E-3		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowercase														=			self.get_dummy_components()
            lowercase														=			StableDiffusionPanoramaPipeline(**__lowerCAmelCase		)
            lowercase														=			sd_pipe.to(__lowerCAmelCase		)
            sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

            lowercase														=			self.get_dummy_inputs(__lowerCAmelCase		)
            lowercase														=			"""french fries"""
            lowercase														=			sd_pipe(**__lowerCAmelCase	, negative_prompt=__lowerCAmelCase		)
            lowercase														=			output.images
            lowercase														=			image[0, -3:, -3:, -1]

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

            lowercase														=			np.array([0.6_1_8_7, 0.5_3_7_5, 0.4_9_1_5, 0.4_1_3_6, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_6, 0.4_7_5_7]		)

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

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowercase														=			self.get_dummy_components()
            lowercase														=			StableDiffusionPanoramaPipeline(**__lowerCAmelCase		)
            lowercase														=			sd_pipe.to(__lowerCAmelCase		)
            sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

            lowercase														=			self.get_dummy_inputs(__lowerCAmelCase		)
            lowercase														=			sd_pipe(**__lowerCAmelCase	, view_batch_size=2		)
            lowercase														=			output.images
            lowercase														=			image[0, -3:, -3:, -1]

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

            lowercase														=			np.array([0.6_1_8_7, 0.5_3_7_5, 0.4_9_1_5, 0.4_1_3_6, 0.4_1_1_4, 0.4_5_6_3, 0.5_1_2_8, 0.4_9_7_6, 0.4_7_5_7]		)

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

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowercase														=			self.get_dummy_components()
            lowercase														=			EulerAncestralDiscreteScheduler(
                beta_start=0.0_0_0_8_5	, beta_end=0.0_1_2	, beta_schedule="""scaled_linear"""		)
            lowercase														=			StableDiffusionPanoramaPipeline(**__lowerCAmelCase		)
            lowercase														=			sd_pipe.to(__lowerCAmelCase		)
            sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

            lowercase														=			self.get_dummy_inputs(__lowerCAmelCase		)
            lowercase														=			sd_pipe(**__lowerCAmelCase		).images
            lowercase														=			image[0, -3:, -3:, -1]

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

            lowercase														=			np.array([0.4_0_2_4, 0.6_5_1_0, 0.4_9_0_1, 0.5_3_7_8, 0.5_8_1_3, 0.5_6_2_2, 0.4_7_9_5, 0.4_4_6_7, 0.4_9_5_2]		)

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







      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            lowercase														=			self.get_dummy_components()
            lowercase														=			PNDMScheduler(
                beta_start=0.0_0_0_8_5	, beta_end=0.0_1_2	, beta_schedule="""scaled_linear"""	, skip_prk_steps=__lowerCAmelCase		)
            lowercase														=			StableDiffusionPanoramaPipeline(**__lowerCAmelCase		)
            lowercase														=			sd_pipe.to(__lowerCAmelCase		)
            sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase		)

            lowercase														=			self.get_dummy_inputs(__lowerCAmelCase		)
            lowercase														=			sd_pipe(**__lowerCAmelCase		).images
            lowercase														=			image[0, -3:, -3:, -1]

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

            lowercase														=			np.array([0.6_3_9_1, 0.6_2_9_1, 0.4_8_6_1, 0.5_1_3_4, 0.5_5_5_2, 0.4_5_7_8, 0.5_0_3_2, 0.5_0_2_3, 0.4_5_3_9]		)

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






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

      def        A__ (		self		):







            """simple docstring"""



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

      def        A__ (		self	, __lowerCAmelCase=0		):







            """simple docstring"""



            lowercase														=			torch.manual_seed(__lowerCAmelCase		)
            lowercase														=			{
                """prompt""": """a photo of the dolomites""",
                """generator""": generator,
                """num_inference_steps""": 3,
                """guidance_scale""": 7.5,
                """output_type""": """numpy""",
            }
            return inputs

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			"""stabilityai/stable-diffusion-2-base"""
            lowercase														=			DDIMScheduler.from_pretrained(__lowerCAmelCase	, subfolder="""scheduler"""		)
            lowercase														=			StableDiffusionPanoramaPipeline.from_pretrained(__lowerCAmelCase	, scheduler=__lowerCAmelCase	, safety_checker=__lowerCAmelCase		)
            pipe.to(__lowerCAmelCase		)
            pipe.set_progress_bar_config(disable=__lowerCAmelCase		)
            pipe.enable_attention_slicing()

            lowercase														=			self.get_inputs()
            lowercase														=			pipe(**__lowerCAmelCase		).images
            lowercase														=			image[0, -3:, -3:, -1].flatten()

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

            lowercase														=			np.array(
                [
                    0.3_6_9_6_8_3_9_2,
                    0.2_7_0_2_5_3_7_2,
                    0.3_2_4_4_6_7_6_6,
                    0.2_8_3_7_9_3_8_7,
                    0.3_6_3_6_3_2_7_4,
                    0.3_0_7_3_3_3_4_7,
                    0.2_7_1_0_0_0_2_7,
                    0.2_7_0_5_4_1_2_5,
                    0.2_5_5_3_6_0_9_6,
                ]		)

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

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			StableDiffusionPanoramaPipeline.from_pretrained(
                """stabilityai/stable-diffusion-2-base"""	, safety_checker=__lowerCAmelCase		)
            lowercase														=			LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
            pipe.to(__lowerCAmelCase		)
            pipe.set_progress_bar_config(disable=__lowerCAmelCase		)
            pipe.enable_attention_slicing()

            lowercase														=			self.get_inputs()
            lowercase														=			pipe(**__lowerCAmelCase		).images
            lowercase														=			image[0, -3:, -3:, -1].flatten()

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

            lowercase														=			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        A__ (		self		):







            """simple docstring"""



            lowercase														=			0

            def callback_fn(__lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		) -> None:
                  lowercase														=			True
                  nonlocal number_of_steps
                  number_of_steps += 1
                  if step == 1:
                        lowercase														=			latents.detach().cpu().numpy()
                        assert latents.shape == (1, 4, 64, 256)
                        lowercase														=			latents[0, -3:, -3:, -1]

                        lowercase														=			np.array(
                            [
                                0.1_8_6_8_1_8_6_9,
                                0.3_3_9_0_7_8_1_6,
                                0.5_3_6_1_2_7_6,
                                0.1_4_4_3_2_8_6_5,
                                -0.0_2_8_5_6_6_1_1,
                                -0.7_3_9_4_1_1_2_3,
                                0.2_3_3_9_7_9_8_7,
                                0.4_7_3_2_2_6_8_2,
                                -0.3_7_8_2_3_1_6_4,
                            ]		)
                        assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5E-2
                  elif step == 2:
                        lowercase														=			latents.detach().cpu().numpy()
                        assert latents.shape == (1, 4, 64, 256)
                        lowercase														=			latents[0, -3:, -3:, -1]

                        lowercase														=			np.array(
                            [
                                0.1_8_5_3_9_6_4_5,
                                0.3_3_9_8_7_2_4_8,
                                0.5_3_7_8_5_5_9,
                                0.1_4_4_3_7_1_4_2,
                                -0.0_2_4_5_5_2_6_1,
                                -0.7_3_3_8_3_1_7,
                                0.2_3_9_9_0_7_5_5,
                                0.4_7_3_5_6_2_7_2,
                                -0.3_7_8_6_5_0_5,
                            ]		)

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

            lowercase														=			False

            lowercase														=			"""stabilityai/stable-diffusion-2-base"""
            lowercase														=			DDIMScheduler.from_pretrained(__lowerCAmelCase	, subfolder="""scheduler"""		)
            lowercase														=			StableDiffusionPanoramaPipeline.from_pretrained(__lowerCAmelCase	, scheduler=__lowerCAmelCase	, safety_checker=__lowerCAmelCase		)
            lowercase														=			pipe.to(__lowerCAmelCase		)
            pipe.set_progress_bar_config(disable=__lowerCAmelCase		)
            pipe.enable_attention_slicing()

            lowercase														=			self.get_inputs()
            pipe(**__lowerCAmelCase	, callback=__lowerCAmelCase	, callback_steps=1		)
            assert callback_fn.has_been_called
            assert number_of_steps == 3







      def        A__ (		self		):







            """simple docstring"""



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

            lowercase														=			"""stabilityai/stable-diffusion-2-base"""
            lowercase														=			DDIMScheduler.from_pretrained(__lowerCAmelCase	, subfolder="""scheduler"""		)
            lowercase														=			StableDiffusionPanoramaPipeline.from_pretrained(__lowerCAmelCase	, scheduler=__lowerCAmelCase	, safety_checker=__lowerCAmelCase		)
            lowercase														=			pipe.to(__lowerCAmelCase		)
            pipe.set_progress_bar_config(disable=__lowerCAmelCase		)
            pipe.enable_attention_slicing(1		)
            pipe.enable_sequential_cpu_offload()

            lowercase														=			self.get_inputs()
            lowercase														=			pipe(**__lowerCAmelCase		)

            lowercase														=			torch.cuda.max_memory_allocated()
            # make sure that less than 5.2 GB is allocated
            assert mem_bytes < 5.5 * 10**9




"""simple docstring"""






import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
       from transformers.models.roberta.modeling_flax_roberta import (
           FlaxRobertaForCausalLM,
           FlaxRobertaForMaskedLM,
           FlaxRobertaForMultipleChoice,
           FlaxRobertaForQuestionAnswering,
           FlaxRobertaForSequenceClassification,
           FlaxRobertaForTokenClassification,
           FlaxRobertaModel,
       )





class 				_A		( unittest.TestCase			):

      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.0_2	, __lowerCAmelCase=4	, ):







            """simple docstring"""



            lowercase														=			parent
            lowercase														=			batch_size
            lowercase														=			seq_length
            lowercase														=			is_training
            lowercase														=			use_attention_mask
            lowercase														=			use_token_type_ids
            lowercase														=			use_labels
            lowercase														=			vocab_size
            lowercase														=			hidden_size
            lowercase														=			num_hidden_layers
            lowercase														=			num_attention_heads
            lowercase														=			intermediate_size
            lowercase														=			hidden_act
            lowercase														=			hidden_dropout_prob
            lowercase														=			attention_probs_dropout_prob
            lowercase														=			max_position_embeddings
            lowercase														=			type_vocab_size
            lowercase														=			type_sequence_label_size
            lowercase														=			initializer_range
            lowercase														=			num_choices

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			ids_tensor([self.batch_size, self.seq_length]	, self.vocab_size		)

            lowercase														=			None
            if self.use_attention_mask:
                  lowercase														=			random_attention_mask([self.batch_size, self.seq_length]		)

            lowercase														=			None
            if self.use_token_type_ids:
                  lowercase														=			ids_tensor([self.batch_size, self.seq_length]	, self.type_vocab_size		)

            lowercase														=			RobertaConfig(
                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        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.prepare_config_and_inputs()
            lowercase							,							lowercase							,							lowercase							,							lowercase														=			config_and_inputs
            lowercase														=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
            return config, inputs_dict







      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.prepare_config_and_inputs()
            lowercase							,							lowercase							,							lowercase							,							lowercase														=			config_and_inputs

            lowercase														=			True
            lowercase														=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
            lowercase														=			ids_tensor([self.batch_size, self.seq_length]	, vocab_size=2		)

            return (
                config,
                input_ids,
                token_type_ids,
                encoder_hidden_states,
                encoder_attention_mask,
            )






@require_flax
class 				_A		( lowerCAmelCase						,					unittest.TestCase			):
      snake_case__    :     Optional[int]													=			True

      snake_case__    :     List[str]													=			(
          (
              FlaxRobertaModel,
              FlaxRobertaForCausalLM,
              FlaxRobertaForMaskedLM,
              FlaxRobertaForSequenceClassification,
              FlaxRobertaForTokenClassification,
              FlaxRobertaForMultipleChoice,
              FlaxRobertaForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			FlaxRobertaModelTester(self		)







      @slow
      def        A__ (		self		):







            """simple docstring"""



            for model_class_name in self.all_model_classes:
                  lowercase														=			model_class_name.from_pretrained("""roberta-base"""	, from_pt=__lowerCAmelCase		)
                  lowercase														=			model(np.ones((1, 1)		)		)
                  self.assertIsNotNone(__lowerCAmelCase		)



