


"""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


a__    :							str            =					logging.get_logger(__name__)

enable_full_determinism()




class 			UpperCamelCase_     (			UpperCamelCase      ,  UpperCamelCase      ,  unittest.TestCase):

	"""simple docstring"""





	snake_case__    :  Optional[Any]    =      UNetaDModel
	snake_case__    :  List[str]    =      "sample"




	@property
	def   UpperCAmelCase_				(			self      :    Tuple       )	->     List[str]:
						__SCREAMING_SNAKE_CASE						       =    4
						__SCREAMING_SNAKE_CASE						       =    3
						__SCREAMING_SNAKE_CASE						       =    (3_2, 3_2)

						__SCREAMING_SNAKE_CASE						       =    floats_tensor((batch_size, num_channels) + sizes       ).to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    torch.tensor([1_0]       ).to(UpperCAmelCase__       )

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




	@property
	def   UpperCAmelCase_				(			self      :    Tuple       )	->     Tuple:
						return (3, 3_2, 3_2)




	@property
	def   UpperCAmelCase_				(			self      :    Tuple       )	->     List[Any]:
						return (3, 3_2, 3_2)




	def   UpperCAmelCase_				(			self      :    Tuple       )	->     str:
						__SCREAMING_SNAKE_CASE						       =    {
						    "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,
						}
						__SCREAMING_SNAKE_CASE						       =    self.dummy_input
						return init_dict, inputs_dict




class 			UpperCamelCase_     (			UpperCamelCase      ,  UpperCamelCase      ,  unittest.TestCase):

	"""simple docstring"""





	snake_case__    :  Union[str, Any]    =      UNetaDModel
	snake_case__    :  str    =      "sample"




	@property
	def   UpperCAmelCase_				(			self      :    List[str]       )	->     str:
						__SCREAMING_SNAKE_CASE						       =    4
						__SCREAMING_SNAKE_CASE						       =    4
						__SCREAMING_SNAKE_CASE						       =    (3_2, 3_2)

						__SCREAMING_SNAKE_CASE						       =    floats_tensor((batch_size, num_channels) + sizes       ).to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    torch.tensor([1_0]       ).to(UpperCAmelCase__       )

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




	@property
	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Tuple:
						return (4, 3_2, 3_2)




	@property
	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     List[Any]:
						return (4, 3_2, 3_2)




	def   UpperCAmelCase_				(			self      :    Any       )	->     Dict:
						__SCREAMING_SNAKE_CASE						       =    {
						    "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"),
						}
						__SCREAMING_SNAKE_CASE						       =    self.dummy_input
						return init_dict, inputs_dict




	def   UpperCAmelCase_				(			self      :    Optional[int]       )	->     Tuple:
						__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"      ,      output_loading_info=UpperCAmelCase__       )

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

						model.to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    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      :    List[str]       )	->     Optional[Any]:
						__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"      ,      output_loading_info=UpperCAmelCase__       )
						model.to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    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      :    Tuple       )	->     Optional[int]:
						# by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
						__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"      ,      output_loading_info=UpperCAmelCase__       )
						model_accelerate.to(UpperCAmelCase__       )
						model_accelerate.eval()

						__SCREAMING_SNAKE_CASE						       =    torch.randn(
						    1      ,      model_accelerate.config.in_channels      ,      model_accelerate.config.sample_size      ,      model_accelerate.config.sample_size      ,      generator=torch.manual_seed(0       )      ,      )
						__SCREAMING_SNAKE_CASE						       =    noise.to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    torch.tensor([1_0] * noise.shape[0]       ).to(UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    model_accelerate(UpperCAmelCase__      ,      UpperCAmelCase__       )["sample"]

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

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

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




	def   UpperCAmelCase_				(			self      :    Optional[int]       )	->     List[Any]:
						__SCREAMING_SNAKE_CASE						       =    UNetaDModel.from_pretrained("fusing/unet-ldm-dummy-update"       )
						model.eval()
						model.to(UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    torch.randn(
						    1      ,      model.config.in_channels      ,      model.config.sample_size      ,      model.config.sample_size      ,      generator=torch.manual_seed(0       )      ,      )
						__SCREAMING_SNAKE_CASE						       =    noise.to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    torch.tensor([1_0] * noise.shape[0]       ).to(UpperCAmelCase__       )

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE						       =    model(UpperCAmelCase__      ,      UpperCAmelCase__       ).sample

						__SCREAMING_SNAKE_CASE						       =    output[0, -1, -3:, -3:].flatten().cpu()
						# fmt: off
						__SCREAMING_SNAKE_CASE						       =    torch.tensor([-13.3_258, -20.1_100, -15.9_873, -17.6_617, -23.0_596, -17.9_419, -13.3_675, -16.1_889, -12.3_800]       )
						# fmt: on

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




class 			UpperCamelCase_     (			UpperCamelCase      ,  UpperCamelCase      ,  unittest.TestCase):

	"""simple docstring"""





	snake_case__    :  Any    =      UNetaDModel
	snake_case__    :  str    =      "sample"




	@property
	def   UpperCAmelCase_				(			self      :    List[str]      ,      UpperCAmelCase__      :    Tuple=(3_2, 3_2)       )	->     int:
						__SCREAMING_SNAKE_CASE						       =    4
						__SCREAMING_SNAKE_CASE						       =    3

						__SCREAMING_SNAKE_CASE						       =    floats_tensor((batch_size, num_channels) + sizes       ).to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    torch.tensor(batch_size * [1_0]       ).to(dtype=torch.intaa      ,      device=UpperCAmelCase__       )

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




	@property
	def   UpperCAmelCase_				(			self      :    Union[str, Any]       )	->     Optional[int]:
						return (3, 3_2, 3_2)




	@property
	def   UpperCAmelCase_				(			self      :    Optional[Any]       )	->     str:
						return (3, 3_2, 3_2)




	def   UpperCAmelCase_				(			self      :    Tuple       )	->     Optional[Any]:
						__SCREAMING_SNAKE_CASE						       =    {
						    "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",
						    ],
						}
						__SCREAMING_SNAKE_CASE						       =    self.dummy_input
						return init_dict, inputs_dict




	@slow
	def   UpperCAmelCase_				(			self      :    Optional[int]       )	->     Dict:
						__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"      ,      output_loading_info=UpperCAmelCase__       )
						self.assertIsNotNone(UpperCAmelCase__       )
						self.assertEqual(len(loading_info["missing_keys"]       )      ,      0       )

						model.to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    self.dummy_input
						__SCREAMING_SNAKE_CASE						       =    floats_tensor((4, 3) + (2_5_6, 2_5_6)       ).to(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    noise
						__SCREAMING_SNAKE_CASE						       =    model(**UpperCAmelCase__       )

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




	@slow
	def   UpperCAmelCase_				(			self      :    Any       )	->     Optional[Any]:
						__SCREAMING_SNAKE_CASE						       =    UNetaDModel.from_pretrained("google/ncsnpp-celebahq-256"       )
						model.to(UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    4
						__SCREAMING_SNAKE_CASE						       =    3
						__SCREAMING_SNAKE_CASE						       =    (2_5_6, 2_5_6)

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

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE						       =    model(UpperCAmelCase__      ,      UpperCAmelCase__       ).sample

						__SCREAMING_SNAKE_CASE						       =    output[0, -3:, -3:, -1].flatten().cpu()
						# fmt: off
						__SCREAMING_SNAKE_CASE						       =    torch.tensor([-4_842.8_691, -6_499.6_631, -3_800.1_953, -7_978.2_686, -10_980.7_129, -20_028.8_535, 8_148.2_822, 2_342.2_905, 567.7_608]       )
						# fmt: on

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




	def   UpperCAmelCase_				(			self      :    int       )	->     Any:
						__SCREAMING_SNAKE_CASE						       =    UNetaDModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update"       )
						model.to(UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    4
						__SCREAMING_SNAKE_CASE						       =    3
						__SCREAMING_SNAKE_CASE						       =    (3_2, 3_2)

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

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE						       =    model(UpperCAmelCase__      ,      UpperCAmelCase__       ).sample

						__SCREAMING_SNAKE_CASE						       =    output[0, -3:, -3:, -1].flatten().cpu()
						# fmt: off
						__SCREAMING_SNAKE_CASE						       =    torch.tensor([-0.0_325, -0.0_900, -0.0_869, -0.0_332, -0.0_725, -0.0_270, -0.0_101, 0.0_227, 0.0_256]       )
						# fmt: on

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




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     List[str]:
						# not required for this model
						pass




"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import _LazyModule


a__    :							Optional[int]            =					{'''tokenization_byt5''': ['''ByT5Tokenizer''']}


if TYPE_CHECKING:
		from .tokenization_byta import ByTaTokenizer
else:
		import sys

		a__    :							Any            =					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
