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





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		   =  StableDiffusionPanoramaPipeline
     SCREAMING_SNAKE_CASE_		   =  TEXT_TO_IMAGE_PARAMS
     SCREAMING_SNAKE_CASE_		   =  TEXT_TO_IMAGE_BATCH_PARAMS
     SCREAMING_SNAKE_CASE_		   =  TEXT_TO_IMAGE_IMAGE_PARAMS
     SCREAMING_SNAKE_CASE_		   =  TEXT_TO_IMAGE_IMAGE_PARAMS






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




           """simple docstring"""







           torch.manual_seed(0  )
           UpperCamelCase__    :   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						,						)
           UpperCamelCase__    :   List[Any]					=				DDIMScheduler()
           torch.manual_seed(0  )
           UpperCamelCase__    :   str					=				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__    :   Optional[Any]					=				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__    :   Union[str, Any]					=				CLIPTextModel(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Union[str, Any]					=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

           UpperCamelCase__    :   Union[str, Any]					=				{
               '''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						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=0  )  ->   Tuple:




           """simple docstring"""







           UpperCamelCase__    :   List[Any]					=				torch.manual_seed(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Union[str, Any]					=				{
               '''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  )  ->   Optional[Any]:




           """simple docstring"""







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

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

           UpperCamelCase__    :   Optional[int]					=				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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Tuple:




           """simple docstring"""







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






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




           """simple docstring"""







           UpperCamelCase__    :   str					=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__    :   str					=				self.get_dummy_components()
           UpperCamelCase__    :   int					=				StableDiffusionPanoramaPipeline(**__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Dict					=				sd_pipe.to(__SCREAMING_SNAKE_CASE  )
           sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )

           UpperCamelCase__    :   Dict					=				self.get_dummy_inputs(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Optional[Any]					=				'''french fries'''
           UpperCamelCase__    :   Any					=				sd_pipe(**__SCREAMING_SNAKE_CASE						,						negative_prompt=__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   Optional[int]					=				output.images
           UpperCamelCase__    :   List[Any]					=				image[0, -3:, -3:, -1]

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

           UpperCamelCase__    :   Tuple					=				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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







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

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

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

           UpperCamelCase__    :   Optional[int]					=				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    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




           """simple docstring"""







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

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

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

           UpperCamelCase__    :   Any					=				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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




           """simple docstring"""







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

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

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

           UpperCamelCase__    :   List[str]					=				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 				_lowerCamelCase (      unittest.TestCase ):





     """simple docstring"""







     def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




           """simple docstring"""







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






     def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE=0  )  ->   Tuple:




           """simple docstring"""







           UpperCamelCase__    :   Optional[int]					=				torch.manual_seed(__SCREAMING_SNAKE_CASE  )
           UpperCamelCase__    :   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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




           """simple docstring"""







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

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

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

           UpperCamelCase__    :   Tuple					=				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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







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

           UpperCamelCase__    :   Union[str, Any]					=				self.get_inputs()
           UpperCamelCase__    :   List[str]					=				pipe(**__SCREAMING_SNAKE_CASE  ).images
           UpperCamelCase__    :   Union[str, Any]					=				image[0, -3:, -3:, -1].flatten()

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

           UpperCamelCase__    :   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    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




           """simple docstring"""







           UpperCamelCase__    :   Optional[int]					=				0

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

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

                       UpperCamelCase__    :   Union[str, Any]					=				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__    :   Optional[int]					=				False

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

           UpperCamelCase__    :   List[str]					=				self.get_inputs()
           pipe(**__SCREAMING_SNAKE_CASE						,						callback=__SCREAMING_SNAKE_CASE						,						callback_steps=1  )
           assert callback_fn.has_been_called
           assert number_of_steps == 3






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




           """simple docstring"""







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

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

           UpperCamelCase__    :   Any					=				self.get_inputs()
           UpperCamelCase__    :   int					=				pipe(**__SCREAMING_SNAKE_CASE  )

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




def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ = 5_0_0_0_0_0_0_0		):
      UpperCamelCase__    :   Any					=				set()
      UpperCamelCase__    :   Any					=				int((limit - 2_4) ** (1 / 2)		)

      UpperCamelCase__    :   Dict					=				set(range(3 ,		prime_square_limit + 1 ,		2		)		)
      primes.add(2		)
      for p in range(3 ,		prime_square_limit + 1 ,		2		):
            if p not in primes:
                  continue
            primes.difference_update(set(range(p * p ,		prime_square_limit + 1 ,		UpperCamelCase__		)		)		)

      for primea in primes:
            UpperCamelCase__    :   Dict					=				primea * primea
            for primea in primes:
                  UpperCamelCase__    :   str					=				primea * primea * primea
                  if square + cube >= limit - 1_6:
                        break
                  for primea in primes:
                        UpperCamelCase__    :   Dict					=				primea * primea * primea * primea
                        UpperCamelCase__    :   Tuple					=				square + cube + tetr
                        if total >= limit:
                              break
                        ret.add(UpperCamelCase__		)

      return len(UpperCamelCase__		)


if __name__ == "__main__":
    print(F'''{solution() = }''')


