





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 			__UpperCamelCase    (							A__	, A__	, unittest.TestCase    ):
     __A	:  Union[str, Any]							  =   StableDiffusionPanoramaPipeline
     __A	:  Tuple							  =   TEXT_TO_IMAGE_PARAMS
     __A	:  str							  =   TEXT_TO_IMAGE_BATCH_PARAMS
     __A	:  List[Any]							  =   TEXT_TO_IMAGE_IMAGE_PARAMS
     __A	:  Tuple							  =   TEXT_TO_IMAGE_IMAGE_PARAMS





     def 		UpperCamelCase(     self					):
        torch.manual_seed(0					)
        _UpperCAmelCase           =   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           =   DDIMScheduler()
        torch.manual_seed(0					)
        _UpperCAmelCase           =   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           =   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=1000				,							)
        _UpperCAmelCase           =   CLIPTextModel(_UpperCamelCase					)
        _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 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase=0					):
        _UpperCAmelCase           =   torch.manual_seed(_UpperCamelCase					)
        _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 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase           =   self.get_dummy_components()
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline(**_UpperCamelCase					)
        _UpperCAmelCase           =   sd_pipe.to(_UpperCamelCase					)
        sd_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   self.get_dummy_inputs(_UpperCamelCase					)
        _UpperCAmelCase           =   sd_pipe(**_UpperCamelCase					).images
        _UpperCAmelCase           =   image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 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 		UpperCamelCase(     self					):
        super().test_inference_batch_consistent(batch_sizes=[1, 2]					)





     def 		UpperCamelCase(     self					):
        super().test_inference_batch_single_identical(batch_size=2				,							expected_max_diff=3.25e-3					)





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase           =   self.get_dummy_components()
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline(**_UpperCamelCase					)
        _UpperCAmelCase           =   sd_pipe.to(_UpperCamelCase					)
        sd_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   self.get_dummy_inputs(_UpperCamelCase					)
        _UpperCAmelCase           =   '''french fries'''
        _UpperCAmelCase           =   sd_pipe(**_UpperCamelCase				,							negative_prompt=_UpperCamelCase					)
        _UpperCAmelCase           =   output.images
        _UpperCAmelCase           =   image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 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 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase           =   self.get_dummy_components()
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline(**_UpperCamelCase					)
        _UpperCAmelCase           =   sd_pipe.to(_UpperCamelCase					)
        sd_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

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

        assert image.shape == (1, 64, 64, 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 		UpperCamelCase(     self					):
        _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(**_UpperCamelCase					)
        _UpperCAmelCase           =   sd_pipe.to(_UpperCamelCase					)
        sd_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

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

        assert image.shape == (1, 64, 64, 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 		UpperCamelCase(     self					):
        _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=_UpperCamelCase					)
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline(**_UpperCamelCase					)
        _UpperCAmelCase           =   sd_pipe.to(_UpperCamelCase					)
        sd_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

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

        assert image.shape == (1, 64, 64, 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 			__UpperCamelCase    (							unittest.TestCase    ):





     def 		UpperCamelCase(     self					):
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 		UpperCamelCase(     self				,							_UpperCamelCase=0					):
        _UpperCAmelCase           =   torch.manual_seed(_UpperCamelCase					)
        _UpperCAmelCase           =   {
            '''prompt''': '''a photo of the dolomites''',
            '''generator''': generator,
            '''num_inference_steps''': 3,
            '''guidance_scale''': 7.5,
            '''output_type''': '''numpy''',
        }
        return inputs





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   '''stabilityai/stable-diffusion-2-base'''
        _UpperCAmelCase           =   DDIMScheduler.from_pretrained(_UpperCamelCase				,							subfolder='''scheduler'''					)
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline.from_pretrained(_UpperCamelCase				,							scheduler=_UpperCamelCase				,							safety_checker=_UpperCamelCase					)
        pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)
        pipe.enable_attention_slicing()

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

        assert image.shape == (1, 512, 2048, 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 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline.from_pretrained(
            '''stabilityai/stable-diffusion-2-base'''				,							safety_checker=_UpperCamelCase					)
        _UpperCAmelCase           =   LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
        pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)
        pipe.enable_attention_slicing()

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

        assert image.shape == (1, 512, 2048, 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 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   0

        def callback_fn(_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase					) -> None:
           _UpperCAmelCase           =   True
           nonlocal number_of_steps
           number_of_steps += 1
           if step == 1:
              _UpperCAmelCase           =   latents.detach().cpu().numpy()
              assert latents.shape == (1, 4, 64, 256)
              _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, 64, 256)
              _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(_UpperCamelCase				,							subfolder='''scheduler'''					)
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline.from_pretrained(_UpperCamelCase				,							scheduler=_UpperCamelCase				,							safety_checker=_UpperCamelCase					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)
        pipe.enable_attention_slicing()

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





     def 		UpperCamelCase(     self					):
        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(_UpperCamelCase				,							subfolder='''scheduler'''					)
        _UpperCAmelCase           =   StableDiffusionPanoramaPipeline.from_pretrained(_UpperCamelCase				,							scheduler=_UpperCamelCase				,							safety_checker=_UpperCamelCase					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)
        pipe.enable_attention_slicing(1					)
        pipe.enable_sequential_cpu_offload()

        _UpperCAmelCase           =   self.get_inputs()
        _UpperCAmelCase           =   pipe(**_UpperCamelCase					)

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






import baseaa
def      A__       (    SCREAMING_SNAKE_CASE_   :      str						)  ->  bytes:






   """simple docstring"""







   return baseaa.baaencode(string.encode('''utf-8'''						)						)
def      A__       (    SCREAMING_SNAKE_CASE_   :      bytes						)  ->  str:






   """simple docstring"""







   return baseaa.baadecode(SCREAMING_SNAKE_CASE_						).decode('''utf-8'''						)


if __name__ == "__main__":
   UpperCAmelCase_           =	"Hello World!"
   UpperCAmelCase_           =	baseaa_encode(test)
   print(encoded)

   UpperCAmelCase_           =	baseaa_decode(encoded)
   print(decoded)