"""simple docstring"""






import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


__lowerCAmelCase    :       Dict					     =logging.get_logger(__name__)

enable_full_determinism()





class  _A						(       snake_case_  ,    snake_case_  ,    unittest.TestCase	):
      snake_case__			:   Tuple							  =	UNetaDModel
      snake_case__			:   Dict							  =	'sample'

      @property
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	4
            lowercase														=	3
            lowercase														=	(32, 32)

            lowercase														=	floats_tensor((batch_size, num_channels) + sizes						).to(__lowerCAmelCase						)
            lowercase														=	torch.tensor([10]						).to(__lowerCAmelCase						)

            return {"sample": noise, "timestep": time_step}

      @property
      def 				A__       (						self						):







            """simple docstring"""



            return (3, 32, 32)

      @property
      def 				A__       (						self						):







            """simple docstring"""



            return (3, 32, 32)







      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	{
                'block_out_channels': (32, 64),
                'down_block_types': ('DownBlock2D', 'AttnDownBlock2D'),
                'up_block_types': ('AttnUpBlock2D', 'UpBlock2D'),
                'attention_head_dim': 3,
                'out_channels': 3,
                'in_channels': 3,
                'layers_per_block': 2,
                'sample_size': 32,
            }
            lowercase														=	self.dummy_input
            return init_dict, inputs_dict






class  _A						(       snake_case_  ,    snake_case_  ,    unittest.TestCase	):
      snake_case__			:   Optional[Any]							  =	UNetaDModel
      snake_case__			:   Optional[Any]							  =	'sample'

      @property
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	4
            lowercase														=	4
            lowercase														=	(32, 32)

            lowercase														=	floats_tensor((batch_size, num_channels) + sizes						).to(__lowerCAmelCase						)
            lowercase														=	torch.tensor([10]						).to(__lowerCAmelCase						)

            return {"sample": noise, "timestep": time_step}

      @property
      def 				A__       (						self						):







            """simple docstring"""



            return (4, 32, 32)

      @property
      def 				A__       (						self						):







            """simple docstring"""



            return (4, 32, 32)

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	{
                'sample_size': 32,
                'in_channels': 4,
                'out_channels': 4,
                'layers_per_block': 2,
                'block_out_channels': (32, 64),
                'attention_head_dim': 32,
                'down_block_types': ('DownBlock2D', 'DownBlock2D'),
                'up_block_types': ('UpBlock2D', 'UpBlock2D'),
            }
            lowercase														=	self.dummy_input
            return init_dict, inputs_dict

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""		,       output_loading_info=__lowerCAmelCase						)

            self.assertIsNotNone(__lowerCAmelCase						)
            self.assertEqual(len(loading_info["""missing_keys"""]						)		,       0						)

            model.to(__lowerCAmelCase						)
            lowercase														=	model(**self.dummy_input						).sample

            assert image is not None, "Make sure output is not None"

      @unittest.skipIf(torch_device != """cuda"""		,       """This test is supposed to run on GPU"""						)
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""		,       output_loading_info=__lowerCAmelCase						)
            model.to(__lowerCAmelCase						)
            lowercase														=	model(**self.dummy_input						).sample

            assert image is not None, "Make sure output is not None"

      @unittest.skipIf(torch_device != """cuda"""		,       """This test is supposed to run on GPU"""						)
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""		,       output_loading_info=__lowerCAmelCase						)
            model_accelerate.to(__lowerCAmelCase						)
            model_accelerate.eval()

            lowercase														=	torch.randn(
                1		,       model_accelerate.config.in_channels		,       model_accelerate.config.sample_size		,       model_accelerate.config.sample_size		,       generator=torch.manual_seed(0						)		,       )
            lowercase														=	noise.to(__lowerCAmelCase						)
            lowercase														=	torch.tensor([10] * noise.shape[0]						).to(__lowerCAmelCase						)

            lowercase														=	model_accelerate(__lowerCAmelCase		,       __lowerCAmelCase						)['sample']

            # two models don't need to stay in the device at the same time
            del model_accelerate
            torch.cuda.empty_cache()
            gc.collect()

            lowercase														=	UNetaDModel.from_pretrained(
                """fusing/unet-ldm-dummy-update"""		,       output_loading_info=__lowerCAmelCase		,       low_cpu_mem_usage=__lowerCAmelCase						)
            model_normal_load.to(__lowerCAmelCase						)
            model_normal_load.eval()
            lowercase														=	model_normal_load(__lowerCAmelCase		,       __lowerCAmelCase						)['sample']

            assert torch_all_close(__lowerCAmelCase		,       __lowerCAmelCase		,       rtol=1E-3						)







      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""						)
            model.eval()
            model.to(__lowerCAmelCase						)

            lowercase														=	torch.randn(
                1		,       model.config.in_channels		,       model.config.sample_size		,       model.config.sample_size		,       generator=torch.manual_seed(0						)		,       )
            lowercase														=	noise.to(__lowerCAmelCase						)
            lowercase														=	torch.tensor([10] * noise.shape[0]						).to(__lowerCAmelCase						)

            with torch.no_grad():
                  lowercase														=	model(__lowerCAmelCase		,       __lowerCAmelCase						).sample

            lowercase														=	output[0, -1, -3:, -3:].flatten().cpu()
            # fmt: off
            lowercase														=	torch.tensor([-1_3.3_2_5_8, -2_0.1_1_0_0, -1_5.9_8_7_3, -1_7.6_6_1_7, -2_3.0_5_9_6, -1_7.9_4_1_9, -1_3.3_6_7_5, -1_6.1_8_8_9, -1_2.3_8_0_0]						)
            # fmt: on

            self.assertTrue(torch_all_close(__lowerCAmelCase		,       __lowerCAmelCase		,       rtol=1E-3						)						)






class  _A						(       snake_case_  ,    snake_case_  ,    unittest.TestCase	):
      snake_case__			:   List[Any]							  =	UNetaDModel
      snake_case__			:   str							  =	'sample'

      @property
      def 				A__       (						self		,       __lowerCAmelCase=(32, 32)						):







            """simple docstring"""



            lowercase														=	4
            lowercase														=	3

            lowercase														=	floats_tensor((batch_size, num_channels) + sizes						).to(__lowerCAmelCase						)
            lowercase														=	torch.tensor(batch_size * [10]						).to(dtype=torch.intaa		,       device=__lowerCAmelCase						)

            return {"sample": noise, "timestep": time_step}

      @property
      def 				A__       (						self						):







            """simple docstring"""



            return (3, 32, 32)

      @property
      def 				A__       (						self						):







            """simple docstring"""



            return (3, 32, 32)

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	{
                'block_out_channels': [32, 64, 64, 64],
                'in_channels': 3,
                'layers_per_block': 1,
                'out_channels': 3,
                'time_embedding_type': 'fourier',
                'norm_eps': 1E-6,
                'mid_block_scale_factor': math.sqrt(2.0						),
                'norm_num_groups': None,
                'down_block_types': [
                    'SkipDownBlock2D',
                    'AttnSkipDownBlock2D',
                    'SkipDownBlock2D',
                    'SkipDownBlock2D',
                ],
                'up_block_types': [
                    'SkipUpBlock2D',
                    'SkipUpBlock2D',
                    'AttnSkipUpBlock2D',
                    'SkipUpBlock2D',
                ],
            }
            lowercase														=	self.dummy_input
            return init_dict, inputs_dict

      @slow
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""		,       output_loading_info=__lowerCAmelCase						)
            self.assertIsNotNone(__lowerCAmelCase						)
            self.assertEqual(len(loading_info["""missing_keys"""]						)		,       0						)

            model.to(__lowerCAmelCase						)
            lowercase														=	self.dummy_input
            lowercase														=	floats_tensor((4, 3) + (256, 256)						).to(__lowerCAmelCase						)
            lowercase														=	noise
            lowercase														=	model(**__lowerCAmelCase						)

            assert image is not None, "Make sure output is not None"

      @slow
      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""						)
            model.to(__lowerCAmelCase						)

            lowercase														=	4
            lowercase														=	3
            lowercase														=	(256, 256)

            lowercase														=	torch.ones((batch_size, num_channels) + sizes						).to(__lowerCAmelCase						)
            lowercase														=	torch.tensor(batch_size * [1E-4]						).to(__lowerCAmelCase						)

            with torch.no_grad():
                  lowercase														=	model(__lowerCAmelCase		,       __lowerCAmelCase						).sample

            lowercase														=	output[0, -3:, -3:, -1].flatten().cpu()
            # fmt: off
            lowercase														=	torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8]						)
            # fmt: on

            self.assertTrue(torch_all_close(__lowerCAmelCase		,       __lowerCAmelCase		,       rtol=1E-2						)						)

      def 				A__       (						self						):







            """simple docstring"""



            lowercase														=	UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""						)
            model.to(__lowerCAmelCase						)

            lowercase														=	4
            lowercase														=	3
            lowercase														=	(32, 32)

            lowercase														=	torch.ones((batch_size, num_channels) + sizes						).to(__lowerCAmelCase						)
            lowercase														=	torch.tensor(batch_size * [1E-4]						).to(__lowerCAmelCase						)

            with torch.no_grad():
                  lowercase														=	model(__lowerCAmelCase		,       __lowerCAmelCase						).sample

            lowercase														=	output[0, -3:, -3:, -1].flatten().cpu()
            # fmt: off
            lowercase														=	torch.tensor([-0.0_3_2_5, -0.0_9_0_0, -0.0_8_6_9, -0.0_3_3_2, -0.0_7_2_5, -0.0_2_7_0, -0.0_1_0_1, 0.0_2_2_7, 0.0_2_5_6]						)
            # fmt: on

            self.assertTrue(torch_all_close(__lowerCAmelCase		,       __lowerCAmelCase		,       rtol=1E-2						)						)







      def 				A__       (						self						):







            """simple docstring"""



            pass




def    __lowercase     (						lowerCamelCase					:     str	):
 UpperCamelCase_					:						Dict =					0
 for ch in input_str:
  UpperCamelCase_					:						Tuple =					ord(lowerCamelCase	)
  UpperCamelCase_					:						str =					pow(2  ,  lowerCamelCase	)

  # If we already turned on bit for current character's unicode
  if bitmap >> ch_unicode & 1 == 1:
   return False
  bitmap |= ch_bit_index_on
 return True


if __name__ == "__main__":
   import doctest

   doctest.testmod()



