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_ (					__a   ,						__a   ,						unittest.TestCase			):
  '''simple docstring'''


  lowercase		:     int            =				StableDiffusionPanoramaPipeline
  lowercase		:     Tuple            =				TEXT_TO_IMAGE_PARAMS
  lowercase		:     str            =				TEXT_TO_IMAGE_BATCH_PARAMS
  lowercase		:     List[Any]            =				TEXT_TO_IMAGE_IMAGE_PARAMS
  lowercase		:     List[Any]            =				TEXT_TO_IMAGE_IMAGE_PARAMS







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

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







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					str      ,	SCREAMING_SNAKE_CASE__    :					str=0     )    ->				Optional[Any]:
         A						:     List[Any]         =torch.manual_seed(A__     )
         A						:     int         ={
             '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 		SCREAMING_SNAKE_CASE_		(				self    :					Any     )    ->				List[str]:
         A						:     List[str]         ='cpu'  # ensure determinism for the device-dependent torch.Generator
         A						:     List[str]         =self.get_dummy_components()
         A						:     str         =StableDiffusionPanoramaPipeline(**A__     )
         A						:     int         =sd_pipe.to(A__     )
         sd_pipe.set_progress_bar_config(disable=A__     )

         A						:     List[str]         =self.get_dummy_inputs(A__     )
         A						:     str         =sd_pipe(**A__     ).images
         A						:     int         =image[0, -3:, -3:, -1]
         assert image.shape == (1, 64, 64, 3)

         A						:     Union[str, Any]         =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 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]     )    ->				List[Any]:
         super().test_inference_batch_consistent(batch_sizes=[1, 2]     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[Any]     )    ->				Optional[int]:
         super().test_inference_batch_single_identical(batch_size=2      ,	expected_max_diff=3.25e-3     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]     )    ->				Union[str, Any]:
         A						:     int         ='cpu'  # ensure determinism for the device-dependent torch.Generator
         A						:     Optional[Any]         =self.get_dummy_components()
         A						:     List[Any]         =StableDiffusionPanoramaPipeline(**A__     )
         A						:     List[Any]         =sd_pipe.to(A__     )
         sd_pipe.set_progress_bar_config(disable=A__     )

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

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

         A						:     str         =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 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]     )    ->				str:
         A						:     str         ='cpu'  # ensure determinism for the device-dependent torch.Generator
         A						:     List[str]         =self.get_dummy_components()
         A						:     int         =StableDiffusionPanoramaPipeline(**A__     )
         A						:     Any         =sd_pipe.to(A__     )
         sd_pipe.set_progress_bar_config(disable=A__     )

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

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

         A						:     Dict         =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 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[Any]     )    ->				Union[str, Any]:
         A						:     Dict         ='cpu'  # ensure determinism for the device-dependent torch.Generator
         A						:     Any         =self.get_dummy_components()
         A						:     str         =EulerAncestralDiscreteScheduler(
             beta_start=0.0_0_0_8_5      ,	beta_end=0.0_1_2      ,	beta_schedule='scaled_linear'     )
         A						:     Any         =StableDiffusionPanoramaPipeline(**A__     )
         A						:     Dict         =sd_pipe.to(A__     )
         sd_pipe.set_progress_bar_config(disable=A__     )

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

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

         A						:     Optional[int]         =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 		SCREAMING_SNAKE_CASE_		(				self    :					Dict     )    ->				int:
         A						:     Optional[int]         ='cpu'  # ensure determinism for the device-dependent torch.Generator
         A						:     Union[str, Any]         =self.get_dummy_components()
         A						:     Dict         =PNDMScheduler(
             beta_start=0.0_0_0_8_5      ,	beta_end=0.0_1_2      ,	beta_schedule='scaled_linear'      ,	skip_prk_steps=A__     )
         A						:     Tuple         =StableDiffusionPanoramaPipeline(**A__     )
         A						:     List[str]         =sd_pipe.to(A__     )
         sd_pipe.set_progress_bar_config(disable=A__     )

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

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

         A						:     Union[str, Any]         =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  SCREAMING_SNAKE_CASE_ (					unittest.TestCase			):
  '''simple docstring'''


  def 		SCREAMING_SNAKE_CASE_		(				self    :					int     )    ->				Dict:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







  def 		SCREAMING_SNAKE_CASE_		(				self    :					int      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]=0     )    ->				str:
         A						:     Optional[int]         =torch.manual_seed(A__     )
         A						:     List[Any]         ={
             'prompt': 'a photo of the dolomites',
             'generator': generator,
             'num_inference_steps': 3,
             'guidance_scale': 7.5,
             'output_type': 'numpy',
         }
         return inputs







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]     )    ->				Any:
         A						:     List[Any]         ='stabilityai/stable-diffusion-2-base'
         A						:     Union[str, Any]         =DDIMScheduler.from_pretrained(A__      ,	subfolder='scheduler'     )
         A						:     Optional[Any]         =StableDiffusionPanoramaPipeline.from_pretrained(A__      ,	scheduler=A__      ,	safety_checker=A__     )
         pipe.to(A__     )
         pipe.set_progress_bar_config(disable=A__     )
         pipe.enable_attention_slicing()

         A						:     int         =self.get_inputs()
         A						:     str         =pipe(**A__     ).images
         A						:     Optional[int]         =image[0, -3:, -3:, -1].flatten()

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

         A						:     List[Any]         =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 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]     )    ->				Tuple:
         A						:     int         =StableDiffusionPanoramaPipeline.from_pretrained(
             'stabilityai/stable-diffusion-2-base'      ,	safety_checker=A__     )
         A						:     int         =LMSDiscreteScheduler.from_config(pipe.scheduler.config     )
         pipe.to(A__     )
         pipe.set_progress_bar_config(disable=A__     )
         pipe.enable_attention_slicing()

         A						:     List[str]         =self.get_inputs()
         A						:     str         =pipe(**A__     ).images
         A						:     str         =image[0, -3:, -3:, -1].flatten()

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

         A						:     Tuple         =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 		SCREAMING_SNAKE_CASE_		(				self    :					List[Any]     )    ->				Union[str, Any]:
         A						:     Union[str, Any]         =0

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

                       A						:     Tuple         =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:
                       A						:     Dict         =latents.detach().cpu().numpy()
                       assert latents.shape == (1, 4, 64, 2_56)
                       A						:     Optional[Any]         =latents[0, -3:, -3:, -1]

                       A						:     Tuple         =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

         A						:     Any         =False

         A						:     Optional[Any]         ='stabilityai/stable-diffusion-2-base'
         A						:     int         =DDIMScheduler.from_pretrained(A__      ,	subfolder='scheduler'     )
         A						:     List[str]         =StableDiffusionPanoramaPipeline.from_pretrained(A__      ,	scheduler=A__      ,	safety_checker=A__     )
         A						:     Dict         =pipe.to(A__     )
         pipe.set_progress_bar_config(disable=A__     )
         pipe.enable_attention_slicing()

         A						:     str         =self.get_inputs()
         pipe(**A__      ,	callback=A__      ,	callback_steps=1     )
         assert callback_fn.has_been_called
         assert number_of_steps == 3







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]     )    ->				Union[str, Any]:
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         A						:     Union[str, Any]         ='stabilityai/stable-diffusion-2-base'
         A						:     Optional[Any]         =DDIMScheduler.from_pretrained(A__      ,	subfolder='scheduler'     )
         A						:     Any         =StableDiffusionPanoramaPipeline.from_pretrained(A__      ,	scheduler=A__      ,	safety_checker=A__     )
         A						:     Optional[int]         =pipe.to(A__     )
         pipe.set_progress_bar_config(disable=A__     )
         pipe.enable_attention_slicing(1     )
         pipe.enable_sequential_cpu_offload()

         A						:     Union[str, Any]         =self.get_inputs()
         A						:     List[Any]         =pipe(**A__     )

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







_lowercase							:      Dict      							='''0.21.0'''

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
    from .utils import rich






