





"""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_     :						int						=					logging.get_logger(__name__)

enable_full_determinism()

class  _UpperCAmelCase  (   __SCREAMING_SNAKE_CASE			,   __SCREAMING_SNAKE_CASE			,   unittest.TestCase		):



      '''simple docstring'''
      lowercase_      : Union[str, Any]	    =   UNetaDModel
      lowercase_      : Any	    =   """sample"""

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       List[str]										=  4
        A_  :       Optional[Any]										=  3
        A_  :       List[str]										=  (3_2, 3_2)

        A_  :       List[Any]										=  floats_tensor((batch_size, num_channels) + sizes  ).to(snake_case_  )
        A_  :       Tuple										=  torch.tensor([1_0]  ).to(snake_case_  )

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

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        return (3, 3_2, 3_2)

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        return (3, 3_2, 3_2)

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       List[Any]										=  {
            'block_out_channels': (3_2, 6_4),
            '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': 3_2,
        }
        A_  :       List[str]										=  self.dummy_input
        return init_dict, inputs_dict


class  _UpperCAmelCase  (   __SCREAMING_SNAKE_CASE			,   __SCREAMING_SNAKE_CASE			,   unittest.TestCase		):



      '''simple docstring'''
      lowercase_      : Optional[Any]	    =   UNetaDModel
      lowercase_      : Tuple	    =   """sample"""

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  4
        A_  :       Dict										=  4
        A_  :       Any										=  (3_2, 3_2)

        A_  :       Tuple										=  floats_tensor((batch_size, num_channels) + sizes  ).to(snake_case_  )
        A_  :       int										=  torch.tensor([1_0]  ).to(snake_case_  )

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

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        return (4, 3_2, 3_2)

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        return (4, 3_2, 3_2)

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Any										=  {
            'sample_size': 3_2,
            'in_channels': 4,
            'out_channels': 4,
            'layers_per_block': 2,
            'block_out_channels': (3_2, 6_4),
            'attention_head_dim': 3_2,
            'down_block_types': ('DownBlock2D', 'DownBlock2D'),
            'up_block_types': ('UpBlock2D', 'UpBlock2D'),
        }
        A_  :       List[str]										=  self.dummy_input
        return init_dict, inputs_dict

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       str										=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'					,      output_loading_info=snake_case_  )

        self.assertIsNotNone(snake_case_  )
        self.assertEqual(len(loading_info['missing_keys']  )					,      0  )

        model.to(snake_case_  )
        A_  :       Dict										=  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 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       Optional[Any]										=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'					,      output_loading_info=snake_case_  )
        model.to(snake_case_  )
        A_  :       Optional[Any]										=  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 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       Any										=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'					,      output_loading_info=snake_case_  )
        model_accelerate.to(snake_case_  )
        model_accelerate.eval()

        A_  :       str										=  torch.randn(
            1					,      model_accelerate.config.in_channels					,      model_accelerate.config.sample_size					,      model_accelerate.config.sample_size					,      generator=torch.manual_seed(0  )					,      )
        A_  :       Any										=  noise.to(snake_case_  )
        A_  :       Any										=  torch.tensor([1_0] * noise.shape[0]  ).to(snake_case_  )

        A_  :       Optional[int]										=  model_accelerate(snake_case_					,      snake_case_  )['sample']

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

        A_						,       A_  :       Tuple										=  UNetaDModel.from_pretrained(
            'fusing/unet-ldm-dummy-update'					,      output_loading_info=snake_case_					,      low_cpu_mem_usage=snake_case_  )
        model_normal_load.to(snake_case_  )
        model_normal_load.eval()
        A_  :       Tuple										=  model_normal_load(snake_case_					,      snake_case_  )['sample']

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[int]										=  UNetaDModel.from_pretrained('fusing/unet-ldm-dummy-update'  )
        model.eval()
        model.to(snake_case_  )

        A_  :       Optional[int]										=  torch.randn(
            1					,      model.config.in_channels					,      model.config.sample_size					,      model.config.sample_size					,      generator=torch.manual_seed(0  )					,      )
        A_  :       Dict										=  noise.to(snake_case_  )
        A_  :       Union[str, Any]										=  torch.tensor([1_0] * noise.shape[0]  ).to(snake_case_  )

        with torch.no_grad():
          A_  :       Optional[int]										=  model(snake_case_					,      snake_case_  ).sample

        A_  :       Union[str, Any]										=  output[0, -1, -3:, -3:].flatten().cpu()
        # fmt: off
        A_  :       Dict										=  torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]  )
        # fmt: on

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


class  _UpperCAmelCase  (   __SCREAMING_SNAKE_CASE			,   __SCREAMING_SNAKE_CASE			,   unittest.TestCase		):



      '''simple docstring'''
      lowercase_      : int	    =   UNetaDModel
      lowercase_      : str	    =   """sample"""

      @property
      def 					lowerCamelCase_			(     self					,      snake_case_=(3_2, 3_2)  ):

        """simple docstring"""

        A_  :       int										=  4
        A_  :       int										=  3

        A_  :       Tuple										=  floats_tensor((batch_size, num_channels) + sizes  ).to(snake_case_  )
        A_  :       List[str]										=  torch.tensor(batch_size * [1_0]  ).to(dtype=torch.intaa					,      device=snake_case_  )

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

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        return (3, 3_2, 3_2)

      @property
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        return (3, 3_2, 3_2)

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[int]										=  {
            'block_out_channels': [3_2, 6_4, 6_4, 6_4],
            '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',
            ],
        }
        A_  :       int										=  self.dummy_input
        return init_dict, inputs_dict

      @slow
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_						,       A_  :       Any										=  UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'					,      output_loading_info=snake_case_  )
        self.assertIsNotNone(snake_case_  )
        self.assertEqual(len(loading_info['missing_keys']  )					,      0  )

        model.to(snake_case_  )
        A_  :       Dict										=  self.dummy_input
        A_  :       List[Any]										=  floats_tensor((4, 3) + (2_5_6, 2_5_6)  ).to(snake_case_  )
        A_  :       Optional[Any]										=  noise
        A_  :       str										=  model(**snake_case_  )

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

      @slow
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       int										=  UNetaDModel.from_pretrained('google/ncsnpp-celebahq-256'  )
        model.to(snake_case_  )

        A_  :       str										=  4
        A_  :       Optional[int]										=  3
        A_  :       Optional[int]										=  (2_5_6, 2_5_6)

        A_  :       Tuple										=  torch.ones((batch_size, num_channels) + sizes  ).to(snake_case_  )
        A_  :       Optional[int]										=  torch.tensor(batch_size * [1E-4]  ).to(snake_case_  )

        with torch.no_grad():
          A_  :       Any										=  model(snake_case_					,      snake_case_  ).sample

        A_  :       int										=  output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        A_  :       str										=  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(snake_case_					,      snake_case_					,      rtol=1E-2  )  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  UNetaDModel.from_pretrained('fusing/ncsnpp-ffhq-ve-dummy-update'  )
        model.to(snake_case_  )

        A_  :       Optional[int]										=  4
        A_  :       str										=  3
        A_  :       Dict										=  (3_2, 3_2)

        A_  :       Union[str, Any]										=  torch.ones((batch_size, num_channels) + sizes  ).to(snake_case_  )
        A_  :       int										=  torch.tensor(batch_size * [1E-4]  ).to(snake_case_  )

        with torch.no_grad():
          A_  :       Optional[Any]										=  model(snake_case_					,      snake_case_  ).sample

        A_  :       Union[str, Any]										=  output[0, -3:, -3:, -1].flatten().cpu()
        # fmt: off
        A_  :       int										=  torch.tensor([-0.03_25, -0.09_00, -0.08_69, -0.03_32, -0.07_25, -0.02_70, -0.01_01, 0.02_27, 0.02_56]  )
        # fmt: on

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

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        pass

import PIL.Image
import PIL.ImageOps
from packaging import version
from PIL import Image


if version.parse(version.parse(PIL.__version__).base_version) >= version.parse('9.1.0'):
					SCREAMING_SNAKE_CASE							:Any							 =				{
					    'linear': PIL.Image.Resampling.BILINEAR,
					    'bilinear': PIL.Image.Resampling.BILINEAR,
					    'bicubic': PIL.Image.Resampling.BICUBIC,
					    'lanczos': PIL.Image.Resampling.LANCZOS,
					    'nearest': PIL.Image.Resampling.NEAREST,
					}
else:
					SCREAMING_SNAKE_CASE							:int							 =				{
					    'linear': PIL.Image.LINEAR,
					    'bilinear': PIL.Image.BILINEAR,
					    'bicubic': PIL.Image.BICUBIC,
					    'lanczos': PIL.Image.LANCZOS,
					    'nearest': PIL.Image.NEAREST,
					}


def 						UpperCAmelCase			(	a_						)		->	Optional[Any]:







				"""simple docstring"""



				__A   			= (images / 2 + 0.5).clamp(0		,   1						)
				__A   			= images.cpu().permute(0		,   2		,   3		,   1						).float().numpy()
				__A   			= numpy_to_pil(a_						)
				return images


def 						UpperCAmelCase			(	a_						)		->	int:







				"""simple docstring"""



				if images.ndim == 3:
								__A   			= images[None, ...]
				__A   			= (images * 2_5_5).round().astype("uint8"						)
				if images.shape[-1] == 1:
								# special case for grayscale (single channel) images
								__A   			= [Image.fromarray(image.squeeze()		,   mode="L"						) for image in images]
				else:
								__A   			= [Image.fromarray(a_						) for image in images]

				return pil_images
