






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_    (    UpperCAmelCase__		,      UpperCAmelCase__		,      unittest.TestCase							):
   _snake_case        =							StableDiffusionPanoramaPipeline
   _snake_case        =							TEXT_TO_IMAGE_PARAMS
   _snake_case        =							TEXT_TO_IMAGE_BATCH_PARAMS
   _snake_case        =							TEXT_TO_IMAGE_IMAGE_PARAMS
   _snake_case        =							TEXT_TO_IMAGE_IMAGE_PARAMS







   def     SCREAMING_SNAKE_CASE__					(self)			->      List[str]:
      """simple docstring"""






      torch.manual_seed(0)
      __snake_case				:      Any     							= UNetaDConditionModel(
          block_out_channels=(3_2, 6_4)					,   layers_per_block=1					,   sample_size=3_2					,   in_channels=4					,   out_channels=4					,   down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')					,   up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')					,   cross_attention_dim=3_2					,   )
      __snake_case				:      Optional[Any]     							= DDIMScheduler()
      torch.manual_seed(0)
      __snake_case				:      int     							= AutoencoderKL(
          block_out_channels=[3_2, 6_4]					,   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=3_2					,   intermediate_size=3_7					,   layer_norm_eps=1E-05					,   num_attention_heads=4					,   num_hidden_layers=5					,   pad_token_id=1					,   vocab_size=1_0_0_0					,   )
      __snake_case				:      Optional[int]     							= CLIPTextModel(__lowerCAmelCase)
      __snake_case				:      Tuple     							= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

      __snake_case				:      Dict     							= {
          '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					,   __a					,   __a=0)			->      Optional[int]:
      """simple docstring"""






      __snake_case				:      int     							= torch.manual_seed(__lowerCAmelCase)
      __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     SCREAMING_SNAKE_CASE__					(self)			->      str:
      """simple docstring"""






      __snake_case				:      List[str]     							= 'cpu'  # ensure determinism for the device-dependent torch.Generator
      __snake_case				:      Tuple     							= self.get_dummy_components()
      __snake_case				:      str     							= StableDiffusionPanoramaPipeline(**__lowerCAmelCase)
      __snake_case				:      int     							= sd_pipe.to(__lowerCAmelCase)
      sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase)

      __snake_case				:      Optional[Any]     							= self.get_dummy_inputs(__lowerCAmelCase)
      __snake_case				:      Any     							= sd_pipe(**__lowerCAmelCase).images
      __snake_case				:      Union[str, Any]     							= image[0, -3:, -3:, -1]
      assert image.shape == (1, 6_4, 6_4, 3)

      __snake_case				:      Tuple     							= 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     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






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







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[Any]:
      """simple docstring"""






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







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






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

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

      assert image.shape == (1, 6_4, 6_4, 3)

      __snake_case				:      str     							= 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     SCREAMING_SNAKE_CASE__					(self)			->      str:
      """simple docstring"""






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

      __snake_case				:      List[Any]     							= self.get_dummy_inputs(__lowerCAmelCase)
      __snake_case				:      Union[str, Any]     							= sd_pipe(**__lowerCAmelCase					,   view_batch_size=2)
      __snake_case				:      str     							= output.images
      __snake_case				:      List[Any]     							= image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)

      __snake_case				:      Tuple     							= 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     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






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

      __snake_case				:      List[Any]     							= self.get_dummy_inputs(__lowerCAmelCase)
      __snake_case				:      Union[str, Any]     							= sd_pipe(**__lowerCAmelCase).images
      __snake_case				:      List[Any]     							= image[0, -3:, -3:, -1]

      assert image.shape == (1, 6_4, 6_4, 3)

      __snake_case				:      Optional[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     SCREAMING_SNAKE_CASE__					(self)			->      List[Any]:
      """simple docstring"""






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

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

      assert image.shape == (1, 6_4, 6_4, 3)

      __snake_case				:      Dict     							= 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 a_    (    unittest.TestCase							):







   def     SCREAMING_SNAKE_CASE__					(self)			->      Union[str, Any]:
      """simple docstring"""






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







   def     SCREAMING_SNAKE_CASE__					(self					,   __a=0)			->      Tuple:
      """simple docstring"""






      __snake_case				:      List[str]     							= torch.manual_seed(__lowerCAmelCase)
      __snake_case				:      int     							= {
          '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]:
      """simple docstring"""






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

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

      assert image.shape == (1, 5_1_2, 2_0_4_8, 3)

      __snake_case				:      Optional[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     SCREAMING_SNAKE_CASE__					(self)			->      List[Any]:
      """simple docstring"""






      __snake_case				:      int     							= StableDiffusionPanoramaPipeline.from_pretrained(
          'stabilityai/stable-diffusion-2-base'					,   safety_checker=__lowerCAmelCase)
      __snake_case				:      Dict     							= LMSDiscreteScheduler.from_config(pipe.scheduler.config)
      pipe.to(__lowerCAmelCase)
      pipe.set_progress_bar_config(disable=__lowerCAmelCase)
      pipe.enable_attention_slicing()

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

      assert image.shape == (1, 5_1_2, 2_0_4_8, 3)

      __snake_case				:      Optional[Any]     							= 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)			->      int:
      """simple docstring"""






      __snake_case				:      List[Any]     							= 0

      def callback_fn(__a					,   __a					,   __a) -> None:
         __snake_case				:      str     							= True
         nonlocal number_of_steps
         number_of_steps += 1
         if step == 1:
            __snake_case				:      Union[str, Any]     							= latents.detach().cpu().numpy()
            assert latents.shape == (1, 4, 6_4, 2_5_6)
            __snake_case				:      List[Any]     							= latents[0, -3:, -3:, -1]

            __snake_case				:      Tuple     							= 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, 6_4, 2_5_6)
            __snake_case				:      Any     							= 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				:      Dict     							= False

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

      __snake_case				:      List[Any]     							= self.get_inputs()
      pipe(**__lowerCAmelCase					,   callback=__lowerCAmelCase					,   callback_steps=1)
      assert callback_fn.has_been_called
      assert number_of_steps == 3







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[Any]:
      """simple docstring"""






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

      __snake_case				:      Tuple     							= 'stabilityai/stable-diffusion-2-base'
      __snake_case				:      int     							= DDIMScheduler.from_pretrained(__lowerCAmelCase					,   subfolder='scheduler')
      __snake_case				:      Tuple     							= StableDiffusionPanoramaPipeline.from_pretrained(__lowerCAmelCase					,   scheduler=__lowerCAmelCase					,   safety_checker=__lowerCAmelCase)
      __snake_case				:      List[str]     							= pipe.to(__lowerCAmelCase)
      pipe.set_progress_bar_config(disable=__lowerCAmelCase)
      pipe.enable_attention_slicing(1)
      pipe.enable_sequential_cpu_offload()

      __snake_case				:      str     							= self.get_inputs()
      __snake_case				:      Dict     							= pipe(**__lowerCAmelCase)

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







'''simple docstring'''






from functools import lru_cache




@lru_cache
def        _SCREAMING_SNAKE_CASE  (				A  :       int  )    -> int:







   """simple docstring"""
   if num < 0:
      raise ValueError('Number should not be negative.'  )

   return 1 if num in (0, 1) else num * factorial(num - 1  )


if __name__ == "__main__":
      import doctest

      doctest.testmod()