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     ):
   __UpperCAmelCase : int      =			StableDiffusionPanoramaPipeline
   __UpperCAmelCase : str      =			TEXT_TO_IMAGE_PARAMS
   __UpperCAmelCase : Optional[int]      =			TEXT_TO_IMAGE_BATCH_PARAMS
   __UpperCAmelCase : str      =			TEXT_TO_IMAGE_IMAGE_PARAMS
   __UpperCAmelCase : List[str]      =			TEXT_TO_IMAGE_IMAGE_PARAMS


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



         torch.manual_seed(0		)
         UpperCAmelCase__								=  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       ,  )
         UpperCAmelCase__								=  DDIMScheduler()
         torch.manual_seed(0		)
         UpperCAmelCase__								=  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		)
         UpperCAmelCase__								=  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       ,  )
         UpperCAmelCase__								=  CLIPTextModel(lowerCAmelCase__		)
         UpperCAmelCase__								=  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

         UpperCAmelCase__								=  {
             "unet": unet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components


   def 			lowercase_      (self     :						Dict       ,  __UpperCAmelCase     :						Optional[Any]       ,  __UpperCAmelCase     :						Optional[int]=0		)						->				Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  torch.manual_seed(lowerCAmelCase__		)
         UpperCAmelCase__								=  {
             "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     :						int		)						->				Optional[int]:
         """simple docstring"""



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase__								=  self.get_dummy_components()
         UpperCAmelCase__								=  StableDiffusionPanoramaPipeline(**lowerCAmelCase__		)
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

         UpperCAmelCase__								=  self.get_dummy_inputs(lowerCAmelCase__		)
         UpperCAmelCase__								=  sd_pipe(**lowerCAmelCase__		).images
         UpperCAmelCase__								=  image[0, -3:, -3:, -1]
         assert image.shape == (1, 6_4, 6_4, 3)

         UpperCAmelCase__								=  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 			lowercase_      (self     :						Tuple		)						->				Optional[Any]:
         """simple docstring"""



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


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



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


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



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase__								=  self.get_dummy_components()
         UpperCAmelCase__								=  StableDiffusionPanoramaPipeline(**lowerCAmelCase__		)
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

         UpperCAmelCase__								=  self.get_dummy_inputs(lowerCAmelCase__		)
         UpperCAmelCase__								=  "french fries"
         UpperCAmelCase__								=  sd_pipe(**lowerCAmelCase__       ,  negative_prompt=lowerCAmelCase__		)
         UpperCAmelCase__								=  output.images
         UpperCAmelCase__								=  image[0, -3:, -3:, -1]

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

         UpperCAmelCase__								=  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 			lowercase_      (self     :						Dict		)						->				List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase__								=  self.get_dummy_components()
         UpperCAmelCase__								=  StableDiffusionPanoramaPipeline(**lowerCAmelCase__		)
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

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

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

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



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase__								=  self.get_dummy_components()
         UpperCAmelCase__								=  EulerAncestralDiscreteScheduler(
             beta_start=0.00085       ,  beta_end=0.012       ,  beta_schedule="scaled_linear"		)
         UpperCAmelCase__								=  StableDiffusionPanoramaPipeline(**lowerCAmelCase__		)
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

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

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

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



         UpperCAmelCase__								=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         UpperCAmelCase__								=  self.get_dummy_components()
         UpperCAmelCase__								=  PNDMScheduler(
             beta_start=0.00085       ,  beta_end=0.012       ,  beta_schedule="scaled_linear"       ,  skip_prk_steps=lowerCAmelCase__		)
         UpperCAmelCase__								=  StableDiffusionPanoramaPipeline(**lowerCAmelCase__		)
         UpperCAmelCase__								=  sd_pipe.to(lowerCAmelCase__		)
         sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

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

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

         UpperCAmelCase__								=  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      A							(				unittest.TestCase     ):


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



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


   def 			lowercase_      (self     :						Any       ,  __UpperCAmelCase     :						Optional[Any]=0		)						->				Optional[Any]:
         """simple docstring"""



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


   def 			lowercase_      (self     :						str		)						->				str:
         """simple docstring"""



         UpperCAmelCase__								=  "stabilityai/stable-diffusion-2-base"
         UpperCAmelCase__								=  DDIMScheduler.from_pretrained(lowerCAmelCase__       ,  subfolder="scheduler"		)
         UpperCAmelCase__								=  StableDiffusionPanoramaPipeline.from_pretrained(lowerCAmelCase__       ,  scheduler=lowerCAmelCase__       ,  safety_checker=lowerCAmelCase__		)
         pipe.to(lowerCAmelCase__		)
         pipe.set_progress_bar_config(disable=lowerCAmelCase__		)
         pipe.enable_attention_slicing()

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

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

         UpperCAmelCase__								=  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 			lowercase_      (self     :						Tuple		)						->				Tuple:
         """simple docstring"""



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

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

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

         UpperCAmelCase__								=  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     :						Dict		)						->				Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  0

         def callback_fn(__UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						torch.FloatTensor		) -> None:
               UpperCAmelCase__								=  True
               nonlocal number_of_steps
               number_of_steps += 1
               if step == 1:
                     UpperCAmelCase__								=  latents.detach().cpu().numpy()
                     assert latents.shape == (1, 4, 6_4, 2_5_6)
                     UpperCAmelCase__								=  latents[0, -3:, -3:, -1]

                     UpperCAmelCase__								=  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__								=  latents.detach().cpu().numpy()
                     assert latents.shape == (1, 4, 6_4, 2_5_6)
                     UpperCAmelCase__								=  latents[0, -3:, -3:, -1]

                     UpperCAmelCase__								=  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__								=  False

         UpperCAmelCase__								=  "stabilityai/stable-diffusion-2-base"
         UpperCAmelCase__								=  DDIMScheduler.from_pretrained(lowerCAmelCase__       ,  subfolder="scheduler"		)
         UpperCAmelCase__								=  StableDiffusionPanoramaPipeline.from_pretrained(lowerCAmelCase__       ,  scheduler=lowerCAmelCase__       ,  safety_checker=lowerCAmelCase__		)
         UpperCAmelCase__								=  pipe.to(lowerCAmelCase__		)
         pipe.set_progress_bar_config(disable=lowerCAmelCase__		)
         pipe.enable_attention_slicing()

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





   def 			lowercase_      (self     :						Dict		)						->				int:
         """simple docstring"""



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

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

         UpperCAmelCase__								=  self.get_inputs()
         UpperCAmelCase__								=  pipe(**lowerCAmelCase__		)

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






from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
						if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						from ...utils.dummy_torch_and_transformers_objects import UnCLIPImageVariationPipeline, UnCLIPPipeline
else:
						from .pipeline_unclip import UnCLIPPipeline
						from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline
						from .text_proj import UnCLIPTextProjModel






