
'''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        UpperCamelCase__     (			UpperCamelCase_					,   UpperCamelCase_					,   unittest.TestCase):
 UpperCAmelCase__			:			str   =    StableDiffusionPanoramaPipeline
 UpperCAmelCase__			:			Any   =    TEXT_TO_IMAGE_PARAMS
 UpperCAmelCase__			:			List[Any]   =    TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase__			:			Optional[Any]   =    TEXT_TO_IMAGE_IMAGE_PARAMS
 UpperCAmelCase__			:			List[Any]   =    TEXT_TO_IMAGE_IMAGE_PARAMS






 def 		lowercase_  (	self     :Optional[Any]	) ->	Optional[int]:







       '''simple docstring'''



       torch.manual_seed(0	)
       __A						=	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		, )
       __A						=	DDIMScheduler()
       torch.manual_seed(0	)
       __A						=	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	)
       __A						=	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=1_000		, )
       __A						=	CLIPTextModel(UpperCamelCase__	)
       __A						=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

       __A						=	{
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components






 def 		lowercase_  (	self     :List[Any]		, _A     :Tuple		, _A     :int=0	) ->	Any:







       '''simple docstring'''



       __A						=	torch.manual_seed(UpperCamelCase__	)
       __A						=	{
           '''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 		lowercase_  (	self     :List[Any]	) ->	str:







       '''simple docstring'''



       __A						=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __A						=	self.get_dummy_components()
       __A						=	StableDiffusionPanoramaPipeline(**UpperCamelCase__	)
       __A						=	sd_pipe.to(UpperCamelCase__	)
       sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

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

       __A						=	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 		lowercase_  (	self     :Dict	) ->	str:







       '''simple docstring'''



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






 def 		lowercase_  (	self     :Union[str, Any]	) ->	Union[str, Any]:







       '''simple docstring'''



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






 def 		lowercase_  (	self     :List[str]	) ->	List[str]:







       '''simple docstring'''



       __A						=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __A						=	self.get_dummy_components()
       __A						=	StableDiffusionPanoramaPipeline(**UpperCamelCase__	)
       __A						=	sd_pipe.to(UpperCamelCase__	)
       sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

       __A						=	self.get_dummy_inputs(UpperCamelCase__	)
       __A						=	'''french fries'''
       __A						=	sd_pipe(**UpperCamelCase__		, negative_prompt=UpperCamelCase__	)
       __A						=	output.images
       __A						=	image[0, -3:, -3:, -1]

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

       __A						=	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 		lowercase_  (	self     :Tuple	) ->	Dict:







       '''simple docstring'''



       __A						=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __A						=	self.get_dummy_components()
       __A						=	StableDiffusionPanoramaPipeline(**UpperCamelCase__	)
       __A						=	sd_pipe.to(UpperCamelCase__	)
       sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

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

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

       __A						=	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 		lowercase_  (	self     :int	) ->	int:







       '''simple docstring'''



       __A						=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __A						=	self.get_dummy_components()
       __A						=	EulerAncestralDiscreteScheduler(
           beta_start=0.00_085		, beta_end=0.012		, beta_schedule='scaled_linear'	)
       __A						=	StableDiffusionPanoramaPipeline(**UpperCamelCase__	)
       __A						=	sd_pipe.to(UpperCamelCase__	)
       sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

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

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

       __A						=	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 		lowercase_  (	self     :Any	) ->	Union[str, Any]:







       '''simple docstring'''



       __A						=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       __A						=	self.get_dummy_components()
       __A						=	PNDMScheduler(
           beta_start=0.00_085		, beta_end=0.012		, beta_schedule='scaled_linear'		, skip_prk_steps=UpperCamelCase__	)
       __A						=	StableDiffusionPanoramaPipeline(**UpperCamelCase__	)
       __A						=	sd_pipe.to(UpperCamelCase__	)
       sd_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

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

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

       __A						=	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        UpperCamelCase__     (			unittest.TestCase):






 def 		lowercase_  (	self     :Dict	) ->	Any:







       '''simple docstring'''



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






 def 		lowercase_  (	self     :Dict		, _A     :Optional[Any]=0	) ->	List[Any]:







       '''simple docstring'''



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






 def 		lowercase_  (	self     :List[Any]	) ->	List[Any]:







       '''simple docstring'''



       __A						=	'''stabilityai/stable-diffusion-2-base'''
       __A						=	DDIMScheduler.from_pretrained(UpperCamelCase__		, subfolder='scheduler'	)
       __A						=	StableDiffusionPanoramaPipeline.from_pretrained(UpperCamelCase__		, scheduler=UpperCamelCase__		, safety_checker=UpperCamelCase__	)
       pipe.to(UpperCamelCase__	)
       pipe.set_progress_bar_config(disable=UpperCamelCase__	)
       pipe.enable_attention_slicing()

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

       assert image.shape == (1, 512, 2_048, 3)

       __A						=	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 		lowercase_  (	self     :Optional[Any]	) ->	Optional[Any]:







       '''simple docstring'''



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

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

       assert image.shape == (1, 512, 2_048, 3)

       __A						=	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 		lowercase_  (	self     :Optional[Any]	) ->	str:







       '''simple docstring'''



       __A						=	0

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

                   __A						=	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:
                   __A						=	latents.detach().cpu().numpy()
                   assert latents.shape == (1, 4, 64, 256)
                   __A						=	latents[0, -3:, -3:, -1]

                   __A						=	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

       __A						=	False

       __A						=	'''stabilityai/stable-diffusion-2-base'''
       __A						=	DDIMScheduler.from_pretrained(UpperCamelCase__		, subfolder='scheduler'	)
       __A						=	StableDiffusionPanoramaPipeline.from_pretrained(UpperCamelCase__		, scheduler=UpperCamelCase__		, safety_checker=UpperCamelCase__	)
       __A						=	pipe.to(UpperCamelCase__	)
       pipe.set_progress_bar_config(disable=UpperCamelCase__	)
       pipe.enable_attention_slicing()

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






 def 		lowercase_  (	self     :Optional[Any]	) ->	Any:







       '''simple docstring'''



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

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

       __A						=	self.get_inputs()
       __A						=	pipe(**UpperCamelCase__	)

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







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCamelCase		:       Any									=     {'configuration_mbart': ['MBART_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MBartConfig', 'MBartOnnxConfig']}

try:
					if not is_sentencepiece_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__UpperCamelCase		:       List[str]									=     ['MBartTokenizer']

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__UpperCamelCase		:       Dict									=     ['MBartTokenizerFast']

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__UpperCamelCase		:       List[str]									=     [
					    'MBART_PRETRAINED_MODEL_ARCHIVE_LIST',
					    'MBartForCausalLM',
					    'MBartForConditionalGeneration',
					    'MBartForQuestionAnswering',
					    'MBartForSequenceClassification',
					    'MBartModel',
					    'MBartPreTrainedModel',
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__UpperCamelCase		:       str									=     [
					    'TFMBartForConditionalGeneration',
					    'TFMBartModel',
					    'TFMBartPreTrainedModel',
					]

try:
					if not is_flax_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__UpperCamelCase		:       List[str]									=     [
					    'FlaxMBartForConditionalGeneration',
					    'FlaxMBartForQuestionAnswering',
					    'FlaxMBartForSequenceClassification',
					    'FlaxMBartModel',
					    'FlaxMBartPreTrainedModel',
					]


if TYPE_CHECKING:
					from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfig, MBartOnnxConfig

					try:
										if not is_sentencepiece_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_mbart import MBartTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_mbart_fast import MBartTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_mbart import (
										    MBART_PRETRAINED_MODEL_ARCHIVE_LIST,
										    MBartForCausalLM,
										    MBartForConditionalGeneration,
										    MBartForQuestionAnswering,
										    MBartForSequenceClassification,
										    MBartModel,
										    MBartPreTrainedModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel, TFMBartPreTrainedModel

					try:
										if not is_flax_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flax_mbart import (
										    FlaxMBartForConditionalGeneration,
										    FlaxMBartForQuestionAnswering,
										    FlaxMBartForSequenceClassification,
										    FlaxMBartModel,
										    FlaxMBartPreTrainedModel,
										)

else:
					import sys

					__UpperCamelCase		:       int									=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





