

"""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_				:      List[Any]           =			logging.get_logger(__name__)

enable_full_determinism()





class  lowerCamelCase    (A__      ,A__      ,unittest.TestCase      ):
			lowerCamelCase__ :		Optional[int]       		=				UNetaDModel
			lowerCamelCase__ :		Optional[Any]       		=				'sample'

			@property
			def        SCREAMING_SNAKE_CASE					(       self					:    str							) ->     int:
						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        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     Tuple:
						return (3, 3_2, 3_2)

			@property
			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     str:
						return (3, 3_2, 3_2)





			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     Union[str, Any]:
						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  lowerCamelCase    (A__      ,A__      ,unittest.TestCase      ):
			lowerCamelCase__ :		Tuple       		=				UNetaDModel
			lowerCamelCase__ :		List[str]       		=				'sample'

			@property
			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     Optional[int]:
						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        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     Optional[int]:
						return (4, 3_2, 3_2)

			@property
			def        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     int:
						return (4, 3_2, 3_2)

			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     List[Any]:
						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        SCREAMING_SNAKE_CASE					(       self					:    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        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]							) ->     List[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        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     Dict:
						# 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        SCREAMING_SNAKE_CASE					(       self					:    str							) ->     int:
						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  lowerCamelCase    (A__      ,A__      ,unittest.TestCase      ):
			lowerCamelCase__ :		str       		=				UNetaDModel
			lowerCamelCase__ :		int       		=				'sample'

			@property
			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]     ,					__UpperCAmelCase					:    Optional[int]=(3_2, 3_2)							) ->     Optional[Any]:
						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        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     List[Any]:
						return (3, 3_2, 3_2)

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

			def        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     int:
						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        SCREAMING_SNAKE_CASE					(       self					:    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        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     str:
						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        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     str:
						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        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     Optional[Any]:
						# not required for this model
						pass





"""simple docstring"""




from ...configuration_utils import PretrainedConfig
from ...utils import logging


A_				:      List[Any]           =			logging.get_logger(__name__)

A_				:      Optional[int]           =			{
    "facebook/dpr-ctx_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-single-nq-base": (
        "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-ctx_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-multiset-base": (
        "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"
    ),
}





class  lowerCamelCase    (A__      ):
			lowerCamelCase__ :		Optional[int]       		=				'dpr'

			def __init__(       self					:    Optional[int]     ,					__UpperCAmelCase					:    List[str]=3_0_5_2_2     ,					__UpperCAmelCase					:    int=7_6_8     ,					__UpperCAmelCase					:    Optional[Any]=1_2     ,					__UpperCAmelCase					:    Dict=1_2     ,					__UpperCAmelCase					:    Any=3_0_7_2     ,					__UpperCAmelCase					:    int="gelu"     ,					__UpperCAmelCase					:    int=0.1     ,					__UpperCAmelCase					:    Optional[Any]=0.1     ,					__UpperCAmelCase					:    Any=5_1_2     ,					__UpperCAmelCase					:    Union[str, Any]=2     ,					__UpperCAmelCase					:    Any=0.02     ,					__UpperCAmelCase					:    Dict=1e-12     ,					__UpperCAmelCase					:    Optional[Any]=0     ,					__UpperCAmelCase					:    Union[str, Any]="absolute"     ,					__UpperCAmelCase					:    int = 0     ,					**__UpperCAmelCase					:    Dict     ,					) ->     Tuple:
						super().__init__(pad_token_id=__UpperCAmelCase     ,					**__UpperCAmelCase							)

						SCREAMING_SNAKE_CASE__										=   vocab_size
						SCREAMING_SNAKE_CASE__										=   hidden_size
						SCREAMING_SNAKE_CASE__										=   num_hidden_layers
						SCREAMING_SNAKE_CASE__										=   num_attention_heads
						SCREAMING_SNAKE_CASE__										=   hidden_act
						SCREAMING_SNAKE_CASE__										=   intermediate_size
						SCREAMING_SNAKE_CASE__										=   hidden_dropout_prob
						SCREAMING_SNAKE_CASE__										=   attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE__										=   max_position_embeddings
						SCREAMING_SNAKE_CASE__										=   type_vocab_size
						SCREAMING_SNAKE_CASE__										=   initializer_range
						SCREAMING_SNAKE_CASE__										=   layer_norm_eps
						SCREAMING_SNAKE_CASE__										=   projection_dim
						SCREAMING_SNAKE_CASE__										=   position_embedding_type


