

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 		SCREAMING_SNAKE_CASE__       (		SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						unittest.TestCase ):







						"""simple docstring"""




						lowercase	:					int           =				StableDiffusionPanoramaPipeline
						lowercase	:					Union[str, Any]           =				TEXT_TO_IMAGE_PARAMS
						lowercase	:					str           =				TEXT_TO_IMAGE_BATCH_PARAMS
						lowercase	:					Any           =				TEXT_TO_IMAGE_IMAGE_PARAMS
						lowercase	:					str           =				TEXT_TO_IMAGE_IMAGE_PARAMS

						def __lowerCamelCase       (					self     )				->						int:




									'''simple docstring'''


									torch.manual_seed(0     )
									__UpperCamelCase     :     Optional[Any]										=   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			,			)
									__UpperCamelCase     :     Tuple										=   DDIMScheduler()
									torch.manual_seed(0     )
									__UpperCamelCase     :     List[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     )
									__UpperCamelCase     :     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=10_00			,			)
									__UpperCamelCase     :     int										=   CLIPTextModel(_lowercase     )
									__UpperCamelCase     :     Union[str, Any]										=   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 __lowerCamelCase       (					self			,			__UpperCamelCase			,			__UpperCamelCase=0     )				->						Optional[int]:




									'''simple docstring'''


									__UpperCamelCase     :     Union[str, Any]										=   torch.manual_seed(_lowercase     )
									__UpperCamelCase     :     str										=   {
									    "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     )				->						List[str]:




									'''simple docstring'''


									__UpperCamelCase     :     Any										=   "cpu"  # ensure determinism for the device-dependent torch.Generator
									__UpperCamelCase     :     Dict										=   self.get_dummy_components()
									__UpperCamelCase     :     Any										=   StableDiffusionPanoramaPipeline(**_lowercase     )
									__UpperCamelCase     :     List[Any]										=   sd_pipe.to(_lowercase     )
									sd_pipe.set_progress_bar_config(disable=_lowercase     )

									__UpperCamelCase     :     Optional[Any]										=   self.get_dummy_inputs(_lowercase     )
									__UpperCamelCase     :     Union[str, Any]										=   sd_pipe(**_lowercase     ).images
									__UpperCamelCase     :     List[Any]										=   image[0, -3:, -3:, -1]
									assert image.shape == (1, 64, 64, 3)

									__UpperCamelCase     :     List[Any]										=   np.array([0.6186, 0.5374, 0.4915, 0.4135, 0.4114, 0.4563, 0.5128, 0.4977, 0.4757]     )

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

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




									'''simple docstring'''


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

						def __lowerCamelCase       (					self     )				->						Tuple:




									'''simple docstring'''


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

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




									'''simple docstring'''


									__UpperCamelCase     :     Any										=   "cpu"  # ensure determinism for the device-dependent torch.Generator
									__UpperCamelCase     :     Union[str, Any]										=   self.get_dummy_components()
									__UpperCamelCase     :     Tuple										=   StableDiffusionPanoramaPipeline(**_lowercase     )
									__UpperCamelCase     :     Tuple										=   sd_pipe.to(_lowercase     )
									sd_pipe.set_progress_bar_config(disable=_lowercase     )

									__UpperCamelCase     :     List[str]										=   self.get_dummy_inputs(_lowercase     )
									__UpperCamelCase     :     int										=   "french fries"
									__UpperCamelCase     :     Union[str, Any]										=   sd_pipe(**_lowercase			,			negative_prompt=_lowercase     )
									__UpperCamelCase     :     str										=   output.images
									__UpperCamelCase     :     Optional[Any]										=   image[0, -3:, -3:, -1]

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

									__UpperCamelCase     :     Union[str, Any]										=   np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]     )

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

						def __lowerCamelCase       (					self     )				->						Any:




									'''simple docstring'''


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

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

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

									__UpperCamelCase     :     str										=   np.array([0.6187, 0.5375, 0.4915, 0.4136, 0.4114, 0.4563, 0.5128, 0.4976, 0.4757]     )

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

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




									'''simple docstring'''


									__UpperCamelCase     :     Dict										=   "cpu"  # ensure determinism for the device-dependent torch.Generator
									__UpperCamelCase     :     List[str]										=   self.get_dummy_components()
									__UpperCamelCase     :     Any										=   EulerAncestralDiscreteScheduler(
									    beta_start=0.00085			,			beta_end=0.012			,			beta_schedule="scaled_linear"     )
									__UpperCamelCase     :     List[str]										=   StableDiffusionPanoramaPipeline(**_lowercase     )
									__UpperCamelCase     :     List[str]										=   sd_pipe.to(_lowercase     )
									sd_pipe.set_progress_bar_config(disable=_lowercase     )

									__UpperCamelCase     :     List[Any]										=   self.get_dummy_inputs(_lowercase     )
									__UpperCamelCase     :     Optional[Any]										=   sd_pipe(**_lowercase     ).images
									__UpperCamelCase     :     int										=   image[0, -3:, -3:, -1]

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

									__UpperCamelCase     :     Tuple										=   np.array([0.4024, 0.6510, 0.4901, 0.5378, 0.5813, 0.5622, 0.4795, 0.4467, 0.4952]     )

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




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




									'''simple docstring'''


									__UpperCamelCase     :     str										=   "cpu"  # ensure determinism for the device-dependent torch.Generator
									__UpperCamelCase     :     str										=   self.get_dummy_components()
									__UpperCamelCase     :     Optional[int]										=   PNDMScheduler(
									    beta_start=0.00085			,			beta_end=0.012			,			beta_schedule="scaled_linear"			,			skip_prk_steps=_lowercase     )
									__UpperCamelCase     :     Dict										=   StableDiffusionPanoramaPipeline(**_lowercase     )
									__UpperCamelCase     :     Dict										=   sd_pipe.to(_lowercase     )
									sd_pipe.set_progress_bar_config(disable=_lowercase     )

									__UpperCamelCase     :     Dict										=   self.get_dummy_inputs(_lowercase     )
									__UpperCamelCase     :     Optional[int]										=   sd_pipe(**_lowercase     ).images
									__UpperCamelCase     :     str										=   image[0, -3:, -3:, -1]

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

									__UpperCamelCase     :     Optional[int]										=   np.array([0.6391, 0.6291, 0.4861, 0.5134, 0.5552, 0.4578, 0.5032, 0.5023, 0.4539]     )

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







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







						"""simple docstring"""




						def __lowerCamelCase       (					self     )				->						str:




									'''simple docstring'''


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

						def __lowerCamelCase       (					self			,			__UpperCamelCase=0     )				->						Optional[Any]:




									'''simple docstring'''


									__UpperCamelCase     :     Any										=   torch.manual_seed(_lowercase     )
									__UpperCamelCase     :     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     )				->						Union[str, Any]:




									'''simple docstring'''


									__UpperCamelCase     :     List[Any]										=   "stabilityai/stable-diffusion-2-base"
									__UpperCamelCase     :     Dict										=   DDIMScheduler.from_pretrained(_lowercase			,			subfolder="scheduler"     )
									__UpperCamelCase     :     Any										=   StableDiffusionPanoramaPipeline.from_pretrained(_lowercase			,			scheduler=_lowercase			,			safety_checker=_lowercase     )
									pipe.to(_lowercase     )
									pipe.set_progress_bar_config(disable=_lowercase     )
									pipe.enable_attention_slicing()

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

									assert image.shape == (1, 5_12, 20_48, 3)

									__UpperCamelCase     :     List[str]										=   np.array(
									    [
									        0.36968392,
									        0.27025372,
									        0.32446766,
									        0.28379387,
									        0.36363274,
									        0.30733347,
									        0.27100027,
									        0.27054125,
									        0.25536096,
									    ]     )

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

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




									'''simple docstring'''


									__UpperCamelCase     :     List[str]										=   StableDiffusionPanoramaPipeline.from_pretrained(
									    "stabilityai/stable-diffusion-2-base"			,			safety_checker=_lowercase     )
									__UpperCamelCase     :     int										=   LMSDiscreteScheduler.from_config(pipe.scheduler.config     )
									pipe.to(_lowercase     )
									pipe.set_progress_bar_config(disable=_lowercase     )
									pipe.enable_attention_slicing()

									__UpperCamelCase     :     Dict										=   self.get_inputs()
									__UpperCamelCase     :     int										=   pipe(**_lowercase     ).images
									__UpperCamelCase     :     Optional[Any]										=   image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 5_12, 20_48, 3)

									__UpperCamelCase     :     List[str]										=   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     )				->						Any:




									'''simple docstring'''


									__UpperCamelCase     :     Any										=   0

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

															__UpperCamelCase     :     Tuple										=   np.array(
															    [
															        0.18681869,
															        0.33907816,
															        0.5361276,
															        0.14432865,
															        -0.02856611,
															        -0.73941123,
															        0.23397987,
															        0.47322682,
															        -0.37823164,
															    ]     )
															assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 5E-2
												elif step == 2:
															__UpperCamelCase     :     Union[str, Any]										=   latents.detach().cpu().numpy()
															assert latents.shape == (1, 4, 64, 2_56)
															__UpperCamelCase     :     Tuple										=   latents[0, -3:, -3:, -1]

															__UpperCamelCase     :     Tuple										=   np.array(
															    [
															        0.18539645,
															        0.33987248,
															        0.5378559,
															        0.14437142,
															        -0.02455261,
															        -0.7338317,
															        0.23990755,
															        0.47356272,
															        -0.3786505,
															    ]     )

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

									__UpperCamelCase     :     Dict										=   False

									__UpperCamelCase     :     List[Any]										=   "stabilityai/stable-diffusion-2-base"
									__UpperCamelCase     :     List[str]										=   DDIMScheduler.from_pretrained(_lowercase			,			subfolder="scheduler"     )
									__UpperCamelCase     :     Optional[Any]										=   StableDiffusionPanoramaPipeline.from_pretrained(_lowercase			,			scheduler=_lowercase			,			safety_checker=_lowercase     )
									__UpperCamelCase     :     Any										=   pipe.to(_lowercase     )
									pipe.set_progress_bar_config(disable=_lowercase     )
									pipe.enable_attention_slicing()

									__UpperCamelCase     :     Tuple										=   self.get_inputs()
									pipe(**_lowercase			,			callback=_lowercase			,			callback_steps=1     )
									assert callback_fn.has_been_called
									assert number_of_steps == 3




						def __lowerCamelCase       (					self     )				->						Tuple:




									'''simple docstring'''


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

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

									__UpperCamelCase     :     Any										=   self.get_inputs()
									__UpperCamelCase     :     Optional[int]										=   pipe(**_lowercase     )

									__UpperCamelCase     :     int										=   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 DummyObject, requires_backends

class       _a  ( metaclass=SCREAMING_SNAKE_CASE__):
			__magic_name__     					=     ["""keras_nlp"""]



			def __init__(			self							:				int   ,	*_lowercase							:				Optional[int]   ,	**_lowercase							:				Dict		)							->					List[Any]:
									requires_backends(self   ,	["keras_nlp"]		)



