





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 					__lowerCAmelCase       (			lowercase_,     lowercase_,     unittest.TestCase     ):
			lowerCamelCase_     :     Optional[int]       			=      StableDiffusionPanoramaPipeline
			lowerCamelCase_     :     Optional[int]       			=      TEXT_TO_IMAGE_PARAMS
			lowerCamelCase_     :     Union[str, Any]       			=      TEXT_TO_IMAGE_BATCH_PARAMS
			lowerCamelCase_     :     Optional[Any]       			=      TEXT_TO_IMAGE_IMAGE_PARAMS
			lowerCamelCase_     :     List[str]       			=      TEXT_TO_IMAGE_IMAGE_PARAMS






			def 			lowerCamelCase  (self    )     ->   Any:


							'''simple docstring'''



							torch.manual_seed(0    )
							snake_case_						:    Tuple				=  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					,							)
							snake_case_						:    List[Any]				=  DDIMScheduler()
							torch.manual_seed(0    )
							snake_case_						:    Any				=  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    )
							snake_case_						:    str				=  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					,							)
							snake_case_						:    Tuple				=  CLIPTextModel(__magic_name__    )
							snake_case_						:    int				=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

							snake_case_						:    Optional[int]				=  {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}
							return components






			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__=0    )     ->   List[str]:


							'''simple docstring'''



							snake_case_						:    str				=  torch.manual_seed(__magic_name__    )
							snake_case_						:    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 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							snake_case_						:    Optional[Any]				=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_						:    List[str]				=  self.get_dummy_components()
							snake_case_						:    Optional[Any]				=  StableDiffusionPanoramaPipeline(**__magic_name__    )
							snake_case_						:    Optional[Any]				=  sd_pipe.to(__magic_name__    )
							sd_pipe.set_progress_bar_config(disable=__magic_name__    )

							snake_case_						:    List[str]				=  self.get_dummy_inputs(__magic_name__    )
							snake_case_						:    Dict				=  sd_pipe(**__magic_name__    ).images
							snake_case_						:    Any				=  image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							snake_case_						:    List[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 			lowerCamelCase  (self    )     ->   List[str]:


							'''simple docstring'''



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






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


							'''simple docstring'''



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






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


							'''simple docstring'''



							snake_case_						:    Any				=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_						:    Optional[Any]				=  self.get_dummy_components()
							snake_case_						:    List[Any]				=  StableDiffusionPanoramaPipeline(**__magic_name__    )
							snake_case_						:    int				=  sd_pipe.to(__magic_name__    )
							sd_pipe.set_progress_bar_config(disable=__magic_name__    )

							snake_case_						:    int				=  self.get_dummy_inputs(__magic_name__    )
							snake_case_						:    Dict				=  '''french fries'''
							snake_case_						:    Tuple				=  sd_pipe(**__magic_name__					,							negative_prompt=__magic_name__    )
							snake_case_						:    Union[str, Any]				=  output.images
							snake_case_						:    Tuple				=  image[0, -3:, -3:, -1]

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

							snake_case_						:    Union[str, Any]				=  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 			lowerCamelCase  (self    )     ->   Tuple:


							'''simple docstring'''



							snake_case_						:    List[str]				=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_						:    Optional[Any]				=  self.get_dummy_components()
							snake_case_						:    Tuple				=  StableDiffusionPanoramaPipeline(**__magic_name__    )
							snake_case_						:    Dict				=  sd_pipe.to(__magic_name__    )
							sd_pipe.set_progress_bar_config(disable=__magic_name__    )

							snake_case_						:    List[str]				=  self.get_dummy_inputs(__magic_name__    )
							snake_case_						:    int				=  sd_pipe(**__magic_name__					,							view_batch_size=2    )
							snake_case_						:    Optional[Any]				=  output.images
							snake_case_						:    Tuple				=  image[0, -3:, -3:, -1]

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

							snake_case_						:    Union[str, Any]				=  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 			lowerCamelCase  (self    )     ->   Tuple:


							'''simple docstring'''



							snake_case_						:    Any				=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							snake_case_						:    Union[str, Any]				=  self.get_dummy_components()
							snake_case_						:    List[Any]				=  EulerAncestralDiscreteScheduler(
							    beta_start=0.00_085					,							beta_end=0.012					,							beta_schedule='''scaled_linear'''    )
							snake_case_						:    Dict				=  StableDiffusionPanoramaPipeline(**__magic_name__    )
							snake_case_						:    str				=  sd_pipe.to(__magic_name__    )
							sd_pipe.set_progress_bar_config(disable=__magic_name__    )

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

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

							snake_case_						:    List[Any]				=  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 			lowerCamelCase  (self    )     ->   Tuple:


							'''simple docstring'''



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

							snake_case_						:    Optional[Any]				=  self.get_dummy_inputs(__magic_name__    )
							snake_case_						:    Dict				=  sd_pipe(**__magic_name__    ).images
							snake_case_						:    List[str]				=  image[0, -3:, -3:, -1]

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

							snake_case_						:    List[Any]				=  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 					__lowerCAmelCase       (			unittest.TestCase     ):






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


							'''simple docstring'''



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






			def 			lowerCamelCase  (self					,							__magic_name__=0    )     ->   Union[str, Any]:


							'''simple docstring'''



							snake_case_						:    Union[str, Any]				=  torch.manual_seed(__magic_name__    )
							snake_case_						:    Union[str, Any]				=  {
							    '''prompt''': '''a photo of the dolomites''',
							    '''generator''': generator,
							    '''num_inference_steps''': 3,
							    '''guidance_scale''': 7.5,
							    '''output_type''': '''numpy''',
							}
							return inputs






			def 			lowerCamelCase  (self    )     ->   Optional[int]:


							'''simple docstring'''



							snake_case_						:    Tuple				=  '''stabilityai/stable-diffusion-2-base'''
							snake_case_						:    Any				=  DDIMScheduler.from_pretrained(__magic_name__					,							subfolder='''scheduler'''    )
							snake_case_						:    List[Any]				=  StableDiffusionPanoramaPipeline.from_pretrained(__magic_name__					,							scheduler=__magic_name__					,							safety_checker=__magic_name__    )
							pipe.to(__magic_name__    )
							pipe.set_progress_bar_config(disable=__magic_name__    )
							pipe.enable_attention_slicing()

							snake_case_						:    Dict				=  self.get_inputs()
							snake_case_						:    Any				=  pipe(**__magic_name__    ).images
							snake_case_						:    Any				=  image[0, -3:, -3:, -1].flatten()

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

							snake_case_						:    List[Any]				=  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 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							snake_case_						:    Any				=  StableDiffusionPanoramaPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-base'''					,							safety_checker=__magic_name__    )
							snake_case_						:    List[str]				=  LMSDiscreteScheduler.from_config(pipe.scheduler.config    )
							pipe.to(__magic_name__    )
							pipe.set_progress_bar_config(disable=__magic_name__    )
							pipe.enable_attention_slicing()

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

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

							snake_case_						:    Optional[int]				=  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 			lowerCamelCase  (self    )     ->   int:


							'''simple docstring'''



							snake_case_						:    List[str]				=  0

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

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

															snake_case_						:    Any				=  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

							snake_case_						:    str				=  False

							snake_case_						:    Tuple				=  '''stabilityai/stable-diffusion-2-base'''
							snake_case_						:    Union[str, Any]				=  DDIMScheduler.from_pretrained(__magic_name__					,							subfolder='''scheduler'''    )
							snake_case_						:    Any				=  StableDiffusionPanoramaPipeline.from_pretrained(__magic_name__					,							scheduler=__magic_name__					,							safety_checker=__magic_name__    )
							snake_case_						:    Tuple				=  pipe.to(__magic_name__    )
							pipe.set_progress_bar_config(disable=__magic_name__    )
							pipe.enable_attention_slicing()

							snake_case_						:    int				=  self.get_inputs()
							pipe(**__magic_name__					,							callback=__magic_name__					,							callback_steps=1    )
							assert callback_fn.has_been_called
							assert number_of_steps == 3






			def 			lowerCamelCase  (self    )     ->   str:


							'''simple docstring'''



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

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

							snake_case_						:    int				=  self.get_inputs()
							snake_case_						:    Optional[int]				=  pipe(**__magic_name__    )

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




'''simple docstring'''

from ....utils import logging


lowerCamelCase							:   Optional[Any]  						=	logging.get_logger(__name__)
class   __lowerCAmelCase						(lowercase_     ):





      '''simple docstring'''






      def __init__(self    :    Optional[int]   ,						UpperCamelCase    :    Union[str, Any]   ,						UpperCamelCase    :    Optional[Any]=None   ,						UpperCamelCase    :    int=2048 ):

             '''simple docstring'''







             lowercase__      			= config.__dict__
             lowercase__      			= modal_hidden_size
             if num_labels:
                    lowercase__      			= num_labels




