





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

enable_full_determinism()



class    A__  (			__snake_case			,  __snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :Optional[int]		       =  UNetaDModel
      _UpperCAmelCase      :str		       =  'sample'





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   str			      =  4
       UpperCamelCase      :   List[Any]			      =  3
       UpperCamelCase      :   Optional[Any]			      =  (32, 32)

       UpperCamelCase      :   List[str]			      =  floats_tensor((batch_size, num_channels) + sizes  ).to(A_  )
       UpperCamelCase      :   Tuple			      =  torch.tensor([10]  ).to(A_  )

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





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return (3, 32, 32)





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return (3, 32, 32)





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Dict			      =  {
           "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,
       }
       UpperCamelCase      :   int			      =  self.dummy_input
       return init_dict, inputs_dict



class    A__  (			__snake_case			,  __snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :List[Any]		       =  UNetaDModel
      _UpperCAmelCase      :Union[str, Any]		       =  'sample'





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Tuple			      =  4
       UpperCamelCase      :   Optional[Any]			      =  4
       UpperCamelCase      :   List[Any]			      =  (32, 32)

       UpperCamelCase      :   str			      =  floats_tensor((batch_size, num_channels) + sizes  ).to(A_  )
       UpperCamelCase      :   List[Any]			      =  torch.tensor([10]  ).to(A_  )

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





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return (4, 32, 32)





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return (4, 32, 32)





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   List[Any]			      =  {
           "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"),
       }
       UpperCamelCase      :   Union[str, Any]			      =  self.dummy_input
       return init_dict, inputs_dict





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      ,       UpperCamelCase      :   int			      =  UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"			,				output_loading_info=A_  )

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

       model.to(A_  )
       UpperCamelCase      :   int			      =  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 			__UpperCamelCase(					self  ):
       '''simple docstring'''




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




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

       UpperCamelCase      :   List[str]			      =  torch.randn(
           1			,				model_accelerate.config.in_channels			,				model_accelerate.config.sample_size			,				model_accelerate.config.sample_size			,				generator=torch.manual_seed(0  )			,				)
       UpperCamelCase      :   List[Any]			      =  noise.to(A_  )
       UpperCamelCase      :   Optional[int]			      =  torch.tensor([10] * noise.shape[0]  ).to(A_  )

       UpperCamelCase      :   str			      =  model_accelerate(A_			,				A_  )["sample"]

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

       UpperCamelCase      ,       UpperCamelCase      :   str			      =  UNetaDModel.from_pretrained(
           "fusing/unet-ldm-dummy-update"			,				output_loading_info=A_			,				low_cpu_mem_usage=A_  )
       model_normal_load.to(A_  )
       model_normal_load.eval()
       UpperCamelCase      :   List[str]			      =  model_normal_load(A_			,				A_  )["sample"]

       assert torch_all_close(A_			,				A_			,				rtol=1e-3  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"  )
       model.eval()
       model.to(A_  )

       UpperCamelCase      :   Optional[int]			      =  torch.randn(
           1			,				model.config.in_channels			,				model.config.sample_size			,				model.config.sample_size			,				generator=torch.manual_seed(0  )			,				)
       UpperCamelCase      :   Optional[Any]			      =  noise.to(A_  )
       UpperCamelCase      :   Any			      =  torch.tensor([10] * noise.shape[0]  ).to(A_  )

       with torch.no_grad():
        UpperCamelCase      :   str			      =  model(A_			,				A_  ).sample

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



class    A__  (			__snake_case			,  __snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :Tuple		       =  UNetaDModel
      _UpperCAmelCase      :Optional[int]		       =  'sample'





      @property
      def 			__UpperCamelCase(					self			,				A_=(32, 32)  ):
       '''simple docstring'''




       UpperCamelCase      :   int			      =  4
       UpperCamelCase      :   int			      =  3

       UpperCamelCase      :   List[str]			      =  floats_tensor((batch_size, num_channels) + sizes  ).to(A_  )
       UpperCamelCase      :   Union[str, Any]			      =  torch.tensor(batch_size * [10]  ).to(dtype=torch.intaa			,				device=A_  )

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





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return (3, 32, 32)





      @property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return (3, 32, 32)





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[int]			      =  {
           "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",
           ],
       }
       UpperCamelCase      :   List[str]			      =  self.dummy_input
       return init_dict, inputs_dict





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




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

       model.to(A_  )
       UpperCamelCase      :   Dict			      =  self.dummy_input
       UpperCamelCase      :   int			      =  floats_tensor((4, 3) + (256, 256)  ).to(A_  )
       UpperCamelCase      :   Tuple			      =  noise
       UpperCamelCase      :   Optional[int]			      =  model(**A_  )

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





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"  )
       model.to(A_  )

       UpperCamelCase      :   int			      =  4
       UpperCamelCase      :   int			      =  3
       UpperCamelCase      :   List[str]			      =  (256, 256)

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

       with torch.no_grad():
        UpperCamelCase      :   List[str]			      =  model(A_			,				A_  ).sample

       UpperCamelCase      :   Optional[Any]			      =  output[0, -3:, -3:, -1].flatten().cpu()
       # fmt: off
       UpperCamelCase      :   Dict			      =  torch.tensor([-48_42.86_91, -64_99.66_31, -38_00.19_53, -79_78.26_86, -1_09_80.71_29, -2_00_28.85_35, 81_48.28_22, 23_42.29_05, 5_67.76_08]  )
       # fmt: on

       self.assertTrue(torch_all_close(A_			,				A_			,				rtol=1e-2  )  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  UNetaDModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update"  )
       model.to(A_  )

       UpperCamelCase      :   Tuple			      =  4
       UpperCamelCase      :   List[Any]			      =  3
       UpperCamelCase      :   Any			      =  (32, 32)

       UpperCamelCase      :   Optional[Any]			      =  torch.ones((batch_size, num_channels) + sizes  ).to(A_  )
       UpperCamelCase      :   List[str]			      =  torch.tensor(batch_size * [1e-4]  ).to(A_  )

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

       UpperCamelCase      :   List[Any]			      =  output[0, -3:, -3:, -1].flatten().cpu()
       # fmt: off
       UpperCamelCase      :   Union[str, Any]			      =  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(A_			,				A_			,				rtol=1e-2  )  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass







import json
import os
import tempfile

from transformers.testing_utils import check_json_file_has_correct_format



class    A__  :
      _UpperCAmelCase      :Union[str, Any]		       =  None





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   int			      =  self.feature_extraction_class(**self.feat_extract_dict  )
       UpperCamelCase      :   Union[str, Any]			      =  json.loads(feat_extract.to_json_string()  )
       for key, value in self.feat_extract_dict.items():
        self.assertEqual(obj[key]			,				A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   str			      =  self.feature_extraction_class(**self.feat_extract_dict  )

       with tempfile.TemporaryDirectory() as tmpdirname:
        UpperCamelCase      :   List[Any]			      =  os.path.join(A_			,				"feat_extract.json"  )
        feat_extract_first.to_json_file(A_  )
        UpperCamelCase      :   str			      =  self.feature_extraction_class.from_json_file(A_  )

       self.assertEqual(feat_extract_second.to_dict()			,				feat_extract_first.to_dict()  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  self.feature_extraction_class(**self.feat_extract_dict  )

       with tempfile.TemporaryDirectory() as tmpdirname:
        UpperCamelCase      :   int			      =  feat_extract_first.save_pretrained(A_  )[0]
        check_json_file_has_correct_format(A_  )
        UpperCamelCase      :   Optional[Any]			      =  self.feature_extraction_class.from_pretrained(A_  )

       self.assertEqual(feat_extract_second.to_dict()			,				feat_extract_first.to_dict()  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Any			      =  self.feature_extraction_class()
       self.assertIsNotNone(A_  )
